Store nn.Parameter in entropy_models.py in nn.ParameterList

[mmuckley]

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]

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]

Another possibility:

    perm = torch.tensor(
        [1, 0, *range(2, x.ndim)], dtype=torch.long, device=x.device
    )

[fracape]

LGTM, might boil down to the same thing under the hood :)

[mmuckley]

Hello @YodaEmbedding, thanks for the suggestion. If it's okay, I would like to argue for the current implementation, as the modification relies on a Python-level range iterator and unpacking it into a list, which can lead to a lot of Python calls under the hood.

When working with frameworks like torch.jit and torch.compile, I often find that these kinds of constructs can be difficult for the compiler, as the most shaky parts of those libraries are around understanding Python. By keeping everything as PyTorch calls, the compilers seem to perform more stably.

[YodaEmbedding]

Sounds good.

[YodaEmbedding]

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]

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.