Add test on number of params
@@ -206,6 +206,43 @@ def test_atcnet(input_sizes):
check_forward_pass(model, input_sizes,)
+def test_atcnet_n_params():
+ """Make sure the number of parameters is the same as in the paper when
+ using the same architecture hyperparameters.
+ """
+ n_windows = 5
+ att_head_dim = 8
+ att_num_heads = 2
+
+ model = ATCNet(
+ n_channels=22,
+ n_classes=4,
+ input_size_s=4.5,
+ sfreq=250,
+ n_windows=n_windows,
+ att_head_dim=att_head_dim,
+ att_num_heads=att_num_heads,
+ )
+
+ n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+ # The paper states the models has around "115.2 K" parameters in its
+ # conclusion. By analyzing the official tensorflow code, we found indeed
+ # 115,172 parameters, but these take into account untrainable batch norm
+ # params, while the number of trainable parameters is 113,732.
+ official_code_nparams = 113_732
+
+ # In our implementation, we need to create the linear embedding mapping
+ # ourselves, in addition to the "square" mappings created by
+ # torch.nn.MultiHeadAttention module. So we need to remove for exceeding
+ # parameters from torch.nn.MultiHeadAttention in our count:
+ embed_dim = att_head_dim * att_num_heads
+ attention_correction = n_windows * (embed_dim * embed_dim + embed_dim) * 4
+ corrected_nparams = n_params - attention_correction
+
+ assert corrected_nparams == official_code_nparams
+
+
@pytest.mark.parametrize(
"n_channels,sfreq,n_classes,input_size_s",
[(20, 128, 5, 30), (10, 256, 4, 20), (1, 64, 2, 30)],
