Loading marblenet model without nemo dependencies (need help with residual block)

[pehonnet]

Hello,

I am trying to load the weights of a model I have trained with nemo, but using only pytorch and torchaudio (as it includes the mfcc extractor as part of the full model). The model is based on Jasper blocks, but I am facing an issue for the blocks with residual connections.

If I use the existing JasperBlock and ConvASREncoder implementation, I am getting this for the first block with residual=true:

    (1): JasperBlock(
      (mconv): ModuleList(
        (0): MaskedConv1d(
          (conv): Conv1d(128, 128, kernel_size=(13,), stride=(1,), padding=(6,), groups=128, bias=False)
        )
        (1): MaskedConv1d(
          (conv): Conv1d(128, 64, kernel_size=(1,), stride=(1,), bias=False)
        )
        (2): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
        (3): ReLU(inplace=True)
        (4): Dropout(p=0.0, inplace=False)
        (5): MaskedConv1d(
          (conv): Conv1d(64, 64, kernel_size=(13,), stride=(1,), padding=(6,), groups=64, bias=False)
        )
        (6): MaskedConv1d(
          (conv): Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)
        )
        (7): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
      )
      (res): ModuleList(
        (0): ModuleList(
          (0): MaskedConv1d(
            (conv): Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)
          )
          (1): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
        )
        (1): ModuleList(
          (0): MaskedConv1d(
            (conv): Conv1d(128, 64, kernel_size=(1,), stride=(1,), bias=False)
          )
          (1): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (mout): Sequential(
        (0): ReLU(inplace=True)
        (1): Dropout(p=0.0, inplace=False)
      )
    )

while in the loaded model (using nemo loading function) I am getting this

    (1): JasperBlock(
      (mconv): ModuleList(
        (0): MaskedConv1d(
          (conv): Conv1d(128, 128, kernel_size=(13,), stride=(1,), padding=(6,), groups=128, bias=False)
        )
        (1): MaskedConv1d(
          (conv): Conv1d(128, 64, kernel_size=(1,), stride=(1,), bias=False)
        )
        (2): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
        (3): ReLU(inplace=True)
        (4): Dropout(p=0.0, inplace=False)
        (5): MaskedConv1d(
          (conv): Conv1d(64, 64, kernel_size=(13,), stride=(1,), padding=(6,), groups=64, bias=False)
        )
        (6): MaskedConv1d(
          (conv): Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)
        )
        (7): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
      )
      (res): ModuleList(
        (0): ModuleList(
          (0): MaskedConv1d(
            (conv): Conv1d(128, 64, kernel_size=(1,), stride=(1,), bias=False)
          )
          (1): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (mout): Sequential(
        (0): ReLU(inplace=True)
        (1): Dropout(p=0.0, inplace=False)
      )
    )

so basically the difference is trying to use the version without nemo there is this additional part:

(0): MaskedConv1d(
            (conv): Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)
          )
          (1): BatchNorm1d(64, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
        )

I am able to get the "correct" architecture with a hack on the function which builds JasperBlock. Before this line: https://github.com/NVIDIA/NeMo/blob/c1abc047d1d5962019f7789b20d8092c7e55b690/nemo/collections/asr/parts/submodules/jasper.py#LL810C33-L810C33
I am adding the following:

res_panes = res_panes[len(res_panes)-1:len(res_panes)]

which helps getting the same structure in all the blocks as "expected".
The issue is then, I am able to load the model but the inference fails with this:

RuntimeError: Given groups=1, weight of size [64, 64, 1], expected input[1, 128, 11809] to have 64 channels, but got 128 channels instead

Could you help me understand what is happening and how to use the correct architectures which lets me load the model and do inference with it?

[pehonnet]

My understanding from the way residual connections work for this model (Marblenet, https://arxiv.org/pdf/2010.13886.pdf) is that at the end of the block, we should have input + layers(input), but in the code I see:

if self.res is not None and self.dense_residual:
    return xs + [out], lens

which is a list of tensors, not just one tensor.
Then, because of that, when we are in block 2 for example, we are working on 2 input tensors coming from block 1, not just 1 tensor.
Is it correct?

[pehonnet]

Sorry, my mistake I had dense_residual = False instead of None in my config...
There is no issue in the code.