KeyError: "state_dict" for every waveglow model I've tried about flowtron HOT 4 CLOSED

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?

from flowtron.

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

from flowtron.

#16

from flowtron.

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!"

from flowtron.