Counting parameters in a neural network

Question

Is there a feature to automatically count parameters in a model? In Keras one can use layer.count_params() to get a count of parameters in a layer or model.summary() to get a detailed summary of parameters and layer sizes for the whole net.

Answer 1

As far as I could see I included all layers that have learnable parameters the correct way; if you spot any issues please feel free to leave a comment or correct my answer.

layers with learnable parameters

countparameters[
l_ConvolutionLayer | l_LinearLayer | l_ConstantArrayLayer | 
l_EmbeddingLayer | l_BatchNormalizationLayer | 
l_BasicRecurrentLayer | l_GatedRecurrentLayer | 
l_LongShortTermMemoryLayer] := 
<|NetExtract[l, "Type"] ->
(NetExtract[NetInitialize@l, "Arrays"] //Query[Select[ArrayQ]] //Map[Length@*Flatten])|>

countparameters[l_SpatialTransformationLayer] := 
<|NetExtract[l, "Type"] -> <|"Parameters" -> 6|>|> 

NetChain and NetGraph

countparameters[l_NetChain | l_NetGraph] := l // Normal // Map[countparameters]

other layers

countparameters[_] := Nothing

---

Example:

lenet = NetChain[{
ConvolutionLayer[20, 5], Ramp, PoolingLayer[2, 2],
ConvolutionLayer[50, 5], Ramp, PoolingLayer[2, 2],
FlattenLayer[], 500, Ramp, 10, SoftmaxLayer[]},
"Output" -> NetDecoder[{"Class", Range[0, 9]}],
"Input" -> NetEncoder[{"Image", {28, 28}, "Grayscale"}]
]

lenet // countparameters

{<|"Convolution" -> <|"Weights" -> 500, "Biases" -> 20|>|>, 
 <|"Convolution" -> <|"Weights" -> 25000, "Biases" -> 50|>|>, 
 <|"Linear" -> <|"Weights" -> 400000, "Biases" -> 500|>|>, 
 <|"Linear" -> <|"Weights" -> 5000, "Biases" -> 10|>|>}

---

Prototype for a fancy version for use with NetGraph

Take a modular net build with NetGraph such as

module[n_] := 
NetChain[{ConvolutionLayer[n, 5], Ramp, PoolingLayer[2, 2]}]

lenet = NetGraph[
<|
"conv_1" -> module[20],
"conv_2" -> module[50], 
"flatten" -> FlattenLayer[],
"dense" -> NetChain[{LinearLayer[500], Ramp, LinearLayer[10], SoftmaxLayer[]}]
|>, 
{"conv_1" -> "conv_2" -> "flatten" -> "dense"}, 
"Output" -> NetDecoder[{"Class", Range[0, 9]}], 
"Input" -> NetEncoder[{"Image", {28, 28}, "Grayscale"}]]

then something like

summary[net_NetGraph] := Module[{sum, parameterCounts},
sum[x_Integer] := x;
sum[x : <|(_ -> _Integer) ..|>] := Total@Values@x;
sum[x_List] := x // Total;
sum[_] := 0;

parameterCounts = net // countparameters;

<|"Total" -> (parameterCounts // Map[sum, #, {0, Infinity}] &) , 
"Details" -> (parameterCounts // Map[sum, #, {1, Infinity}] &) |>
]

gives a nice overview how the learnable parameters are distributed throughout the network architecture.

lenet // summary

<|
 "Total" -> 431080, 
 "Details" ->
 <|"conv_1" -> 520, "conv_2" -> 25050, "flatten" -> 0, "dense" -> 405510|>
|>