KeyError: "state_dict" for every waveglow model I've tried

[lachjames]

Hi :)

I've successfully trained a flowtron model and I'm now trying to use it for inference with inference.py. I've tried a few different waveglow versions (v5 in the waveglow repo, v3 from the nvidia NGC catalog), both with and without running convert_model.py on them, and I always get the same error:

Traceback (most recent call last):
File "inference.py", line 132, in
args.id, args.n_frames, args.sigma, args.gate, args.seed)
File "inference.py", line 55, in infer
state_dict = torch.load(flowtron_path, map_location='cpu')["state_dict"]
KeyError: 'state_dict'

Clearly I'm doing something wrong - I was hoping you might be able to help me out.

Thanks :)

EDIT: I think the problem has nothing to do with waveglow at all, but instead is (possibly) that I can't use a flowtron checkpoint for inference. Is this the case?

[adrianastan]

The pretrained Waveglow model files contain the entire model. Use smth like this in the inference.py file:

m = torch.load(flowtron_path, map_location='cpu') state_dict = m['model'].state_dict() model.load_state_dict(state_dict)

Or a more general solution, as in train.py: https://github.com/NVIDIA/flowtron/blob/master/train.py#L82

[lachjames]

Thanks, that helped :) I was able to change line 55 in inference.py to

state_dict = torch.load(flowtron_path, map_location='cpu')["model"].state_dict()

and that allowed me to use the checkpoint.

Perhaps some code which tries the current method and, if the key state_dict is not found, tries this instead, would be a good addition to inference.py?

[rafaelvalle]

#16

[serg06]

Thanks, that helped :) I was able to change line 55 in inference.py to

state_dict = torch.load(flowtron_path, map_location='cpu')["model"].state_dict()

and that allowed me to use the checkpoint.

Perhaps some code which tries the current method and, if the key state_dict is not found, tries this instead, would be a good addition to inference.py?

I was trying to use a Flowtron checkpoint and run into the same issue. Unfortunately your fix didn't work for me:

Traceback (most recent call last):
  File "inference.py", line 140, in <module>
    args.id, args.n_frames, args.sigma, args.gate, args.seed)
  File "inference.py", line 63, in infer
    model.load_state_dict(state_dict)
  File "/home/serguei/anaconda3/envs/flowtron/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1052, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))

RuntimeError: Error(s) in loading state_dict for Flowtron:

        size mismatch for speaker_embedding.weight: copying a param with shape torch.Size([123, 128]) from checkpoint, the shape in current model is torch.Size([1, 128]).

However, I was able to fix it by doing this:

    model_on_disk = torch.load(flowtron_path, map_location='cpu')
    if 'state_dict' in model_on_disk:
        state_dict = model_on_disk['state_dict']
        model.load_state_dict(state_dict)
    elif 'model' in model_on_disk:
        model = model_on_disk['model'].cuda()
    else:
        assert False, "cannot load model!"