InterDigitalInc / CompressAI

A PyTorch library and evaluation platform for end-to-end compression research
https://interdigitalinc.github.io/CompressAI/
BSD 3-Clause Clear License
1.21k stars 232 forks source link

Store `nn.Parameter` in `entropy_models.py` in `nn.ParameterList` #284

Closed mmuckley closed 6 months ago

mmuckley commented 6 months ago

This PR proposes to store parameters in entropy_models.py in an nn.ParameterList instead of its current string-based lookup. The primary reason to do so is to make EntropyBottleneck more friendly for torch.compile, where the current implementation fails to compile for certain backends (in my own experience, dynamo). The primary reason seems to be that the current implementation relies too much on Python strings and class attributes to access the parameters, whereas the new implementation makes this more clear at the PyTorch level, which helps the compiler.

A major drawback to the PR merging would be that it breaks backwards compatibility. I've included some state_dict adjustments that would allow loading old checkpoints, but I understand this may not be ideal. Also, new checkpoints would not be loadable by older versions of compressai.

The PR also includes a compile test for verifying the implementation.

Happy to see this merged or closed, depending on maintainer preference.

fracape commented 6 months ago

Thank you @mmuckley for the PR! Checking a couple of things on our side. I'll probably tag a new version after merging to indicate the potential minor backward compatibility issue re: new checkpoints with older versions of the package. This is a research project, I think it's worth it and totally fine in that direction (new checkpoints - old package).

YodaEmbedding commented 6 months ago

The PR looks good; I tested the current implementation with torch.compile, and it works on my machine with PyTorch 2.3.0.


On a semi-related side note, the current ELIC implementation fails with torch.compile, though perhaps that will resolve itself with future torch versions.

fracape commented 6 months ago

The PR looks good; I tested the current implementation with torch.compile, and it works on my machine with PyTorch 2.3.0.

On a semi-related side note, the current ELIC implementation fails with torch.compile, though perhaps that will resolve itself with future torch versions.

I limited the supported torch <2.3 this week, temporarily. I'm getting slightly different results in eval_model video which breaks CI that compares with expected results produced with earlier versions.