Closed mmuckley closed 4 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).
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.
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.
This PR proposes to store parameters in
entropy_models.py
in annn.ParameterList
instead of its current string-based lookup. The primary reason to do so is to makeEntropyBottleneck
more friendly fortorch.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.