Mismatched keys in pretrained JETS model

[iamanigeeit]

Hello @imdanboy,

When i try to load the pretrained model from HuggingFace, it seems that tts.discriminator.msd.discriminators.0.layers should have weight_v and weight_g but it only has weight.

os.chdir(LJSPEECH_DIR)
pretrained_dir = LJSPEECH_DIR / "exp/tts_train_jets_raw_phn_tacotron_g2p_en_no_space"
pretrained_model_file = pretrained_dir / "train.total_count.ave_5best_new.pth"
pretrained_tts = Text2Speech.from_pretrained(
    train_config=pretrained_dir / "config.yaml",
    model_file=pretrained_model_file,
    device=device
)
pretrained_model = pretrained_tts.model
os.chdir(PWD)

Error:

RuntimeError: Error(s) in loading state_dict for ESPnetTTSModel:
	Missing key(s) in state_dict: "tts.discriminator.msd.discriminators.0.layers.0.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.0.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.1.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.1.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.2.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.2.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.3.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.3.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.4.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.4.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.5.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.5.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.6.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.6.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.7.weight_g", "tts.discriminator.msd.discriminators.0.layers.7.weight_v". 
	Unexpected key(s) in state_dict: "tts.discriminator.msd.discriminators.0.layers.0.0.weight", "tts.discriminator.msd.discriminators.0.layers.1.0.weight", "tts.discriminator.msd.discriminators.0.layers.2.0.weight", "tts.discriminator.msd.discriminators.0.layers.3.0.weight", "tts.discriminator.msd.discriminators.0.layers.4.0.weight", "tts.discriminator.msd.discriminators.0.layers.5.0.weight", "tts.discriminator.msd.discriminators.0.layers.6.0.weight", "tts.discriminator.msd.discriminators.0.layers.7.weight".

I think i can fix this by calculating weight_v and weight_g but this error only happened recently. All other parameters in the pretrained model use weight_v and weight_g.

[iamanigeeit]

Seems like due to the weight norm bug being fixed in commit 1b45bdc

I propose to close #4595

[iamanigeeit]

@kan-bayashi The mismatched keys problem should apply to all HiFiGAN pretrained models. The below code should take the old .pth, compute the weight norm and write to a new .pth (keeping the module order).

import torch
import yaml
from collections import OrderedDict
from espnet2.gan_tts.hifigan.hifigan import HiFiGANMultiScaleDiscriminator as Msd

config_file = 'config.yaml'
pth_file = 'train.total_count.ave_5best.pth'
new_pth_file = 'train.total_count.ave_5best_new.pth'

state_dict = torch.load(pth_file)
msd_prefix = 'tts.discriminator.msd.'
msd_state_dict = {k[len(msd_prefix):]: v for (k, v) in state_dict.items() if k.startswith(msd_prefix)}

with open(config_file, 'r') as f:
    config = yaml.safe_load(f)
discrim_config = config['tts_conf']['discriminator_params']
scale_prefix = 'scale_'
msd_config = {}
for k, v in discrim_config.items():
    if not k.startswith('period'):
        if k.startswith(scale_prefix):
            key = k[len(scale_prefix):]
        else:
            key = k
        msd_config[key] = v

msd_config['discriminator_params']['use_weight_norm'] = False
msd = Msd(**msd_config)

for discrim in msd.discriminators:
    discrim.apply_weight_norm()

new_msd_state = msd.state_dict(prefix=msd_prefix)

before = True
before_keys, after_keys = [], []
for k in state_dict:
    if k.startswith(msd_prefix):
        before = False
        continue
    else:
        if before:
            before_keys.append(k)
        else:
            after_keys.append(k)

new_state_dict = OrderedDict()
for k in before_keys:
    new_state_dict[k] = state_dict[k]
for k in new_msd_state:
    new_state_dict[k] = new_msd_state[k]
for k in after_keys:
    new_state_dict[k] = state_dict[k]

for x in new_state_dict:
    print(x)

torch.save(new_state_dict, new_pth_file)

[imdanboy]

I faced the same problem and had to checkout to old commit (specified at huggingface) 😅

[sw005320]

Thanks for the report.
@kan-bayashi, can you deal with this?

[kan-bayashi]

I will try to fix without re-uploading all of models. Those are too much :(

[kan-bayashi]

Fixed in #5240

[nellorebhanuteja]

@kan-bayashi Looks like this error still occurs when when you pip install espnet

[kan-bayashi]

Not yet released. Please use the master.

[nellorebhanuteja]

I did

pip install git+https://github.com/espnet/espnet

Getting the following error while loading VITS model

RuntimeError: Error(s) in loading state_dict for ESPnetTTSModel:
        Missing key(s) in state_dict: "tts.discriminator.msd.discriminators.0.layers.0.0.weight", "tts.discriminator.msd.discriminators.0.layers.1.0.weight", "tts.discriminator.msd.discriminators.0.layers.2.0.weight", "tts.discriminator.msd.discriminators.0.layers.3.0.weight", "tts.discriminator.msd.discriminators.0.layers.4.0.weight", "tts.discriminator.msd.discriminators.0.layers.5.0.weight", "tts.discriminator.msd.discriminators.0.layers.6.0.weight", "tts.discriminator.msd.discriminators.0.layers.7.weight". 
        Unexpected key(s) in state_dict: "tts.discriminator.msd.discriminators.0.layers.0.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.0.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.1.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.1.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.2.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.2.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.3.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.3.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.4.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.4.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.5.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.5.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.6.0.weight_g", "tts.discriminator.msd.discriminators.0.layers.6.0.weight_v", "tts.discriminator.msd.discriminators.0.layers.7.weight_g", "tts.discriminator.msd.discriminators.0.layers.7.weight_v". 

[kan-bayashi]

It seems your case is slightly different from the above case.
You try to load weights with norm but the model does not use norm.
Maybe you already changed config not to use use_weight_norm or use_spectral_norm.
If not, please provide the reproducible codes.

[nellorebhanuteja]

Right.
I changed the value of use_weight_norm. Issue resolved.