Closed YodaEmbedding closed 7 months ago
base config (ReduceLROnPlateau)?
Yes, default is fine. Judging by this implementation's reported results, it is possible to get very close to the paper's results. Like the paper, they also use 8000 images from ImageNet, but hopefully that has minimal effect.
Hi @YodaEmbedding, was trying to figure out where the problem is, but my results using the above (vimeo90k, default ReduceLROnPlateau) are not the expected performance:
Did you launch on your end or would you have a clue regarding a wrong setup here?
I think that may have been due to a bug in compress/decompress. I fixed it. The current implementation seems to be working now with some limited testing. I think forward
was fine, so it may have trained correctly. Thus, you may be able to get away with just loading the weights to check, and not have to train from scratch.
Summary of recent changes:
cheng2020-anchor-elic
which uses suboptimal architecture settings.elic2022-official
which uses official architecture. The number of channels weren't always specified, so I made some educated guesses.elic2022-chandelier
which uses the exact same architecture as Vincent Chandelier's ELiC-ReImplemetation. He also provides some pretrained checkpoints. To load them with this implementation, use the Python script provided in the top-level comment to rename the weights. The resulting RD performance is identical.LatentCodec
s._forward_twopass
. I originally wrote _forward_twopass_faster
based on the original paper's description. Both seem to work and are nearly equivalent, but I've defaulted to Chandelier's twopass since it seems like it would be very slightly more robust under certain architectures.I plugged in Chandelier's pretrained weights (trained on ImageNet 8000) and then "finetuned" on Vimeo90K for a few epochs.
Chandelier pretrained model finetuned on Vimeo90K for 1 epoch | Chandelier pretrained model finetuned on Vimeo90K for 4 epochs |
---|---|
Not unexpectedly, "finetuning" on a different, larger dataset actually improves RD performance further.
Not shown: finetuning 0 epochs, since I'm a bit too lazy to rerun things.
Suggestions:
elic2022-chandelier
, elic2022-official
, and cheng2020-anchor-checkerboard
in that order. See original post for suggested hyperparameters. (Ignore.)
Regarding the - means
trick, it wouldn't work with GMM. We can't really pre-compute the distributions with GMM because they depend on too many parameters.
Also, not sure about ELIC but FYI in my experiments Checkerboard + cheng2020 results differ by quite a lot between K=1 and K=3
Checkerboard and ELIC models.
Differences from original paper:
GaussianConditionalLatentCodec
once that PR is accepted. I guess the STE " - means" offset might need some tweaking too once there's more than one mean...Checklist:
++model.name="cheng2020-anchor-checkerboard" ++hp.N=128 ++hp.M=192
++model.name="cheng2020-anchor-checkerboard" ++hp.N=192 ++hp.M=320
++model.name="cheng2020-anchor-checkerboard" ++hp.N=192 ++hp.M=320
++model.name="elic2022-official" ++hp.N=128 ++hp.M=192 ++hp.groups='[16, 16, 32, 64, 64]'
++model.name="elic2022-official" ++hp.N=192 ++hp.M=320 ++hp.groups='[16, 16, 32, 64, 192]'
++model.name="elic2022-official" ++hp.N=192 ++hp.M=320 ++hp.groups='[16, 16, 32, 64, 192]'
++model.name="elic2022-chandelier" ++hp.N=128 ++hp.M=192 ++hp.groups='[16, 16, 32, 64, 64]'
++model.name="elic2022-chandelier" ++hp.N=192 ++hp.M=320 ++hp.groups='[16, 16, 32, 64, 192]'
++model.name="elic2022-chandelier" ++hp.N=192 ++hp.M=320 ++hp.groups='[16, 16, 32, 64, 192]'
results/image/*/compressai-*.json
.Compatibility check with Chandelier ELiC-ReImplemetation pretrained models
Download models from: https://github.com/VincentChandelier/ELiC-ReImplemetation#available-checkpoint
Save the following python script:
Click for python script
```python import argparse import torch FORWARD_MAPPING = { "g_a": "g_a", "g_s": "g_s", "h_a": "latent_codec.hyper.h_a", "h_s": "latent_codec.hyper.h_s", "entropy_bottleneck": "latent_codec.hyper.entropy_bottleneck", "cc_transforms.0": "latent_codec.y.channel_context.y1", "cc_transforms.1": "latent_codec.y.channel_context.y2", "cc_transforms.2": "latent_codec.y.channel_context.y3", "cc_transforms.3": "latent_codec.y.channel_context.y4", "context_prediction.0": "latent_codec.y.latent_codec.y0.context_prediction", "context_prediction.1": "latent_codec.y.latent_codec.y1.context_prediction", "context_prediction.2": "latent_codec.y.latent_codec.y2.context_prediction", "context_prediction.3": "latent_codec.y.latent_codec.y3.context_prediction", "context_prediction.4": "latent_codec.y.latent_codec.y4.context_prediction", "ParamAggregation.0": "latent_codec.y.latent_codec.y0.entropy_parameters", "ParamAggregation.1": "latent_codec.y.latent_codec.y1.entropy_parameters", "ParamAggregation.2": "latent_codec.y.latent_codec.y2.entropy_parameters", "ParamAggregation.3": "latent_codec.y.latent_codec.y3.entropy_parameters", "ParamAggregation.4": "latent_codec.y.latent_codec.y4.entropy_parameters", } REVERSE_MAPPING = { "latent_codec.y.latent_codec.y0.y.gaussian_conditional": "gaussian_conditional", "latent_codec.y.latent_codec.y1.y.gaussian_conditional": "gaussian_conditional", "latent_codec.y.latent_codec.y2.y.gaussian_conditional": "gaussian_conditional", "latent_codec.y.latent_codec.y3.y.gaussian_conditional": "gaussian_conditional", "latent_codec.y.latent_codec.y4.y.gaussian_conditional": "gaussian_conditional", } def _rename_key(key): found = False for src_prefix, dst_prefix in FORWARD_MAPPING.items(): if key.startswith(src_prefix): new_key = f"{dst_prefix}{key[len(src_prefix):]}" yield new_key found = True for dst_prefix, src_prefix in REVERSE_MAPPING.items(): if key.startswith(src_prefix): new_key = f"{dst_prefix}{key[len(src_prefix):]}" yield new_key found = True if found: return raise RuntimeError(f"Unmapped key: {key}") def rename_keys(state_dict): max_len = max(len(key) for key in state_dict.keys()) new_state_dict = {} for key, value in state_dict.items(): for new_key in _rename_key(key): print(f"{key:<{max_len}} -> {new_key:<{max_len}}") new_state_dict[new_key] = value return new_state_dict def build_parser(): parser = argparse.ArgumentParser() parser.add_argument( "--input", type=str, required=True, help="Path to the weights file" ) parser.add_argument( "--output", type=str, required=True, help="Path to the output file" ) return parser def main(): parser = build_parser() args = parser.parse_args() state_dict = torch.load(args.input) state_dict = rename_keys(state_dict) torch.save(state_dict, args.output) if __name__ == "__main__": main() ```Then, run:
Then load the model checkpoint in CompressAI Trainer: