YodaEmbedding
commented
2 years ago There are two possible methods to handle this.
-
Run
update_modelto modify the checkpoint:python -m compressai.utils.update_model checkpoint.pth.tarThis also freezes the checkpoint, removes some state (e.g. optimizer), and adds a hash to the filename. If that is not desired, the alternative is...
-
After loading the model, call
net.update(force=True):net = models_video[args.model](quality=3) net.update(force=True)
update_model is necessary to update the parameters needed for entropy coding. These parameters are not used during training, so they can be added after training. More info here: https://github.com/InterDigitalInc/CompressAI/issues/5#issuecomment-724881519
Personally, I've created a wrapper that injects net.update(force=True) into compressai.utils.eval_model.
# Filename: personal_code/utils/compressai/update_and_eval_model.py
import compressai.utils.eval_model.__main__ as eval_model_main
import torch
import torch.nn as nn
from compressai.zoo import load_state_dict
from compressai.zoo.image import model_architectures as architectures
def load_checkpoint(arch: str, checkpoint_path: str) -> nn.Module:
ckpt = torch.load(checkpoint_path)
state_dict = ckpt["state_dict"] if "state_dict" in ckpt else ckpt
state_dict = load_state_dict(state_dict)
model = architectures[arch].from_state_dict(state_dict).eval()
model.update(force=True)
return model
def main(argv):
eval_model_main.__dict__["load_checkpoint"] = load_checkpoint
eval_model_main.main(argv)# Filename: personal_code/utils/compressai/__main__.py
import importlib
import sys
if __name__ == "__main__":
_, util_name, *argv = sys.argv
if util_name == "update_and_eval_model":
from . import update_and_eval_model
main = update_and_eval_model.main
else:
module = importlib.import_module(
f"compressai.utils.{util_name}.__main__"
)
main = module.main
main(argv)Usage:
python -m personal_code.utils.compressai update_and_eval_model checkpoint $DATASET_PATH --verbose -a=bmshj2018-factorized -p checkpoint.pth.tar > results.json
amarzullo24
chyomin06
Bug
I am trying to load a checkpoint after training the ssf2020 model. However, when I load the checkpoint, a size mismatch error occurs
To Reproduce
Run the following script (provided you already have a "checkpoint.pth.tar")
You should get the following error:
Expected behavior
The model state loads properly.
Environment
Please copy and paste the output from
python3 -m torch.utils.collect_envAdditional context
I have noticed that if I print the actual value of the parameters in the stack trace, the corresponding tensor in the state_dict is actually empty:
update
As expected, by commenting lines 439-482 of compressai.models.video.google.py I get the loading work. However, I guess this is not a proper solution.