Count Model Parameters

mediumPROModelingPyTorch ModelingPytorch

Problem Statement

Count the total number of trainable parameters in a neural network.

Background

Understanding model size is important for deployment and memory planning. Each layer's parameters include weights and biases, and the total determines the model's capacity and memory footprint.

Your Task

Write a function count_parameters(in_features, hidden_features, out_features) that builds a two-layer network (with ReLU) and returns a breakdown of its trainable parameters.

Output Format

Return a dictionary with keys: "total_params" (int), "layer1_params" (int), "layer2_params" (int).

Example

Input

in_features=10, hidden_features=20, out_features=5

Output

{'total_params': 325, 'layer1_params': 220, 'layer2_params': 105}

Reasoning

Inputs: in_features=10 is the input vector size, hidden_features=20 is the width of the hidden layer, and out_features=5 is the number of outputs. The network maps $10 \to 20 \to 5$ , with ReLU in between.
Layer 1: a linear map from 10 inputs to 20 outputs holds a $20 \times 10$ weight matrix plus one bias per output, so $10 \cdot 20 + 20 = 200 + 20 = 220$ parameters.
Layer 2: the same rule from 20 inputs to 5 outputs gives a $5 \times 20$ weight matrix plus 5 biases: $20 \cdot 5 + 5 = 100 + 5 = 105$ .
ReLU is a fixed elementwise function $\max(0, x)$ with nothing to learn, so it contributes 0.
Output: the total is $220 + 105 = 325$ , giving {'total_params': 325, 'layer1_params': 220, 'layer2_params': 105}. All three values are plain integers.
Notice that the hidden width appears in both layers, so the count grows roughly as $(\text{in} + \text{out}) \times \text{hidden}$ : for $784 \to 128 \to 10$ the first layer alone is $100480$ .

Constraints

Linear layer params = in_features * out_features + out_features (bias)
Use param.numel() to count elements

Learn this in

PyTorch FundamentalsFoundations Language Model BackboneVLM Study Plan

More from Pytorch

Apply Softmax for Predictionseasy Build with nn.Sequentialeasy Use nn.Flatten Layereasy Create a Simple Neural Network Classeasy Compose Multiple Transformsmedium

Concept primerPREVIEW

1. Background Knowledge

Neural network parameters are the learnable weights and biases in layers that are optimized during training. In PyTorch, nn.Linear(in_features, out_features) creates a layer with:

Weight matrix: \mathbf{W} $\in \mathbb{R}^{$ \text{out_features} $\times$ \text{in_features} $}$ → \text{in_features} $\times$ \text{out_features} parameters
Bias vector: \mathbf{b} $\in \mathbb{R}^{$ \text{out_features} $}$ → \text{out_features} parameters
Total: \text{in_features} $\times$ \text{out_features} $+$ \text{out_features}

Trainable parameters have requires_grad=True (default for nn.Linear). ReLU activations have 0 parameters as they are fixed functions: $f(x) = \max(0, x)$ .

Key PyTorch methods:

for name, param in model.named_parameters():
 if param.requires_grad:
 count = param.numel() # Total elements in tensor

Sample calculation: For Linear(10→20) + Linear(20→5):

Layer 1: $10 \times 20 + 20 = 220$
Layer 2: $20 \times 5 + 5 = 105$
Total: 325 trainable parameters

2. Algorithm Approach

Direct iteration over parameters (most reliable):

1. Build model: Sequential(Linear1 → ReLU → Linear2)
2. Iterate model.named_parameters()
3. For each param: if requires_grad, add param.numel()
4. Track layer-specific counts using name filtering

Mathematical formulation: Let $L_1 =$ \text{Linear} $(n_{in}, n_h)$ , $L_2 =$ \text{Linear} $(n_h, n_{out})$

params(L_i) = n_{prev} × n_i + n_i
total_params = params(L_1) + params(L_2)

Alternative: sum(p.numel() for p in model.parameters() if p.requires_grad)

3. Step-by-Step Strategy

def count_parameters(in_features, hidden_features, out_features):
 # Step 1: Create model
 model = nn.Sequential(
 nn.Linear(in_features, hidden_features),
 nn.ReLU(),
 nn.Linear(hidden_features, out_features)
 )

 # Step 2: Initialize counters
 total = 0
 layer1 = 0
 layer2 = 0

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results