KeyError: 'g_a.0.weight' in Model update

Hi all,

I am trying to finetune a pre-trained model on my dataset. The packages version I used are listed as follows:

PyTorch: Version: 1.8.1+cu111   
CompressAI: Version: 1.1.4

I mainly modified one line in examples/train.py and changed it to finetune.py. The line I modified changed to the following:

net = models[args.model](quality=args.quality, metric="mse", pretrained=True)

And I used the following script to finetune the pretrained model and save it to a checkpoint.

python finetune.py -m cheng2020-anchor -q 1 -d ../../test-pictures-finetuned --save --cuda

And it successfully generates the checkpoints I need. Then, I follow the steps provided in the tutorial. To update the CDF, and export the checkpoint. I did it with the following command.

python -m compressai.utils.update_model --architecture mean-scale-hyperprior -n cheng2020-anchor-finetuned -d ../updated-checkpoints checkpoint_best_loss.pth.tar

And the error occurred:

Traceback (most recent call last):                                                                                                                                                       
  File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main                                                                                                                   
    return _run_code(code, main_globals, None,                                                                                                                                           
  File "/usr/lib/python3.8/runpy.py", line 87, in _run_code                                                                                                                              
    exec(code, run_globals)                                                                                                                                                              
  File "/usr/local/lib/python3.8/dist-packages/compressai/utils/update_model/__main__.py", line 139, in <module>                                                                         
    main(sys.argv[1:])                                                                                                                                                                   
  File "/usr/local/lib/python3.8/dist-packages/compressai/utils/update_model/__main__.py", line 110, in main                                                                             
    net = model_cls.from_state_dict(state_dict)                                                                                                                                          
  File "/usr/local/lib/python3.8/dist-packages/compressai/models/priors.py", line 274, in from_state_dict                                                                                
    N = state_dict["g_a.0.weight"].size(0)                                                                                                                                               
KeyError: 'g_a.0.weight' 

I changed the --architecture to all other architectures, {factorized-prior,jarhp,mean-scale-hyperprior,scale-hyperprior} . None of these works. Could you help me with this? Thanks a lot!

[jbegaint]

Great 😃.

[jsw-zorro]

Hi Jean @jbegaint , Never mind. I found I just need to regenerate the checkpoint and then it can be updated using the updated pip package. Thanks for the help.