RuntimeError when loading model checkpoint

Hi @andi611 @leo19941227

When I am trying to load the model state dict for my pytorch model, I am getting the following error:

  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1406, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for CustomModel:
        size mismatch for extractor_s3prl.extracter._melscale.fb: copying a param with shape torch.Size([513, 64]) from checkpoint, the shape in current model is torch.Size([0]).

Can you guide me on how to resolve the same?

Below is just a snippet roughly representing my model code for your reference. After training my model, I saved its model state.

class CustomModel(nn.Module):
    def __init__(self,  ...):

        super(CustomModel, self).__init__()

        self.s3prl_feature_type = 'baseline_local'
        self.extractor_s3prl =  torch.hub.load('s3prl/s3prl', self.s3prl_feature_type, model_config='path_to_custom_mel.yaml').to('cuda')

Sincerely, Soham Tiwari

[andi611]

@leo19941227 I found this part in the source code:

if n_stft is None or n_stft == 0:
    warnings.warn(
        'Initialization of torchaudio.transforms.MelScale with an unset weight '
        '`n_stft=None` is deprecated and will be removed in release 0.10. '
        'Please set a proper `n_stft` value. Typically this is `n_fft // 2 + 1`. '
        'Refer to https://github.com/pytorch/audio/issues/1510 '
        'for more details.'

Should we pass an n_stft value when we initialize MelScale?

[leo19941227]

Hey actually this is torchaudio’s bug and will be resolved after torchaudio==v0.10.0 You can refer to this issue and I provide a workaround below with adding just one line:

import torch
import torch.nn as nn
from s3prl.hub import baseline_local

class CustomModel(nn.Module):
    def __init__(self):

        super(CustomModel, self).__init__()

        self.extractor_s3prl = baseline_local(model_config="mel.yaml").to("cuda") 

    def forward(self, x):

        wavs = [wav.to('cuda') for wav in x.squeeze(1)]
        out_s3prl = self.extractor_s3prl(wavs)['last_hidden_state']
        x = torch.unsqueeze(out_s3prl, 1)

        output_dict = {
            'x': x
        }

        return output_dict
    
model = CustomModel()
x = torch.rand([1, 192000]).to('cuda')
out = model(x)
checkpoint = {
    'model': model.state_dict()
}
torch.save(checkpoint, 'custom_model.pth')

# loading state dict
model = CustomModel()

# ADD THIS LINE
out = model(x)

checkpoint = torch.load('custom_model.pth')
model.load_state_dict(checkpoint['model'])