Calculate Number of Parameters in Neural Network

Problem

Given a neural network architecture specified as a list of layer dimensions, calculate the total number of trainable parameters (weights and biases) in the network.

Solution

def count_parameters(layer_dims: list[int], use_bias: bool = True) -> int:
    total = 0
    for i in range(len(layer_dims) - 1):
        fan_in = layer_dims[i]
        fan_out = layer_dims[i + 1]
        total += fan_in * fan_out
        if use_bias:
            total += fan_out
    return total

Explanation

1. For each pair of consecutive layers with dimensions (fan_in, fan_out), the weight matrix has fan_in * fan_out parameters.
2. If biases are used, each neuron in the output layer has one bias parameter, adding fan_out parameters.
3. Sum across all layers to get the total parameter count.
4. Example: [784, 128, 64, 10] has 784128 + 128 + 12864 + 64 + 64*10 + 10 = 109,386 parameters.

Complexity

• Time: O(L) where L is the number of layers
• Space: O(1)