How to use newly trained model

[Franckyboi]

To add a new model, the readme states to add the small_vqvae, small_prior... etc into MODELS to make_models.py But how exactly? should these be relative paths/absolute paths/ the path to the checkpoint files/ "small_prior", or is there something left to do with the checkpoint files?

[prafullasd]

You want to add an entry for restore_prior="path/to/checkpoint" inside the corresponding hps dictionary, in your case "small_prior". For example, if your checkpoint is stored at "/root/jukebox/logs/small_prior/checkpoint_latest.pth.tar" and you're going to replace the top_level with "small_prior, you want the entry in hparams.py to be:

small_prior = Hyperparams(
    n_ctx=8192,
    ... other hps here ...
    init_scale=0.7,
    restore_prior="/root/jukebox/logs/small_prior/checkpoint_latest.pth.tar", 
)

If you trained with labels/lyrics, make sure you also set these hps correctly: y_bins, max_bow_genre_size for the genre/artist information min_duration, max_duration, t_bins for the timing information use_tokens, n_tokens, n_vocab for the lyric information

For example, if you trained with --hps=small_single_enc_dec_prior --labels --y_bins=(10,100) --max_bow_genre_size=1 --min_duration=60 --max_duration=600 --t_bins=64 --use_tokens=True --n_tokens=384 --n_vocab=79 then you want to set

small_single_enc_dec_prior = Hyperparams(
    n_ctx=6144,
    ... other hps here ...
    single_enc_dec=True,
    labels=True,
    y_bins=(10,100),
    max_bow_genre_size=1,
    .... other hps from training here...
    n_tokens=384,
    n_vocab=79,
)

so that when you restore the model for sampling they're correctly set.

[made-by-chris]

To piggyback tangentially, could you theoretically train locally and then load it into the colab environment and use it in a hosted set up? I can run the sampler locally.. but maybe i can train. Also - has anyone managed to apply the training command within with colab env?

[Franckyboi]

You want to add an entry for restore_prior="path/to/checkpoint" inside the corresponding hps dictionary, in your case "small_prior". For example, if your checkpoint is stored at "/root/jukebox/logs/small_prior/checkpoint_latest.pth.tar" and you're going to replace the top_level with "small_prior,

Thank you, I also had to manually set the param "labels" to False in at least one of the newly trained model parts:

small_prior = Hyperparams(
    n_ctx=8192,
    ... other hps here ...
    init_scale=0.7,
    labels=False
    restore_prior="/root/jukebox/logs/small_prior/checkpoint_latest.pth.tar", 
)

[Franckyboi]

I trained the models with the standard commands in the readme file but I think the way the models are stored differs from how they are opened to use in make_models.py. I'm getting the following errors:

File "/home/user/Downloads/jukebox/jukebox/make_models.py", line 60, in restore
    model.load_state_dict(checkpoint['model'])
File "/home/user/anaconda3/envs/jukebox/lib/python3.7/site-packages/torch/nn/modules/module.py", line 777, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for SimplePrior:
	Unexpected key(s) in state_dict: "conditioner_blocks.0.cond.model.2.0.blocks.0.model.1.weight", "conditioner_blocks.0.cond.model.2.0.blocks.0.model.1.bias"

followed by some of these:

size mismatch for conditioner_blocks.0.x_emb.weight: copying a param with shape torch.Size([1024, 1024]) from checkpoint, the shape in current model is torch.Size([512, 1024]).

Is this something I can fix with loading the checkpoint files differently or are my models corrupted?

[heewooj]

It's hard to say without looking at the exact commands, but these usually mean that some flags used during training are missing at sampling time

[vigliensoni]

Hi @Superchaschper, were you able to sort out the issue with the unexpected key(s) in state_dict?

[Franckyboi]

Hi @Superchaschper, were you able to sort out the issue with the unexpected key(s) in state_dict?

Hey @vigliensoni, No, I wasn't able to solve it and after getting the same errors in another attempt, I haven't tried it since.