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

compressai/entropy_models/entropy_models.py

@@ -360,20 +360,24 @@ def __init__(
         scale = self.init_scale ** (1 / (len(self.filters) + 1))
         channels = self.channels
 
+        self.matrices = nn.ParameterList()
+        self.biases = nn.ParameterList()
+        self.factors = nn.ParameterList()
+
         for i in range(len(self.filters) + 1):
             init = np.log(np.expm1(1 / scale / filters[i + 1]))
             matrix = torch.Tensor(channels, filters[i + 1], filters[i])
             matrix.data.fill_(init)
-            self.register_parameter(f"_matrix{i:d}", nn.Parameter(matrix))
+            self.matrices.append(nn.Parameter(matrix))
 
             bias = torch.Tensor(channels, filters[i + 1], 1)
             nn.init.uniform_(bias, -0.5, 0.5)
-            self.register_parameter(f"_bias{i:d}", nn.Parameter(bias))
+            self.biases.append(nn.Parameter(bias))
 
             if i < len(self.filters):
                 factor = torch.Tensor(channels, filters[i + 1], 1)
                 nn.init.zeros_(factor)
-                self.register_parameter(f"_factor{i:d}", nn.Parameter(factor))
+                self.factors.append(nn.Parameter(factor))
 
         self.quantiles = nn.Parameter(torch.Tensor(channels, 1, 3))
         init = torch.Tensor([-self.init_scale, 0, self.init_scale])
@@ -433,24 +437,23 @@ def _logits_cumulative(self, inputs: Tensor, stop_gradient: bool) -> Tensor:
         # TorchScript not yet working (nn.Mmodule indexing not supported)
         logits = inputs
         for i in range(len(self.filters) + 1):
-            matrix = getattr(self, f"_matrix{i:d}")
+            matrix = self.matrices[i]
             if stop_gradient:
                 matrix = matrix.detach()
             logits = torch.matmul(F.softplus(matrix), logits)
 
-            bias = getattr(self, f"_bias{i:d}")
+            bias = self.biases[i]
             if stop_gradient:
                 bias = bias.detach()
-            logits += bias
+            logits = logits + bias
 
             if i < len(self.filters):
-                factor = getattr(self, f"_factor{i:d}")
+                factor = self.factors[i]
                 if stop_gradient:
                     factor = factor.detach()
-                logits += torch.tanh(factor) * torch.tanh(logits)
+                logits = logits + torch.tanh(factor) * torch.tanh(logits)
         return logits
 
-    @torch.jit.unused
     def _likelihood(
         self, inputs: Tensor, stop_gradient: bool = False
     ) -> Tuple[Tensor, Tensor, Tensor]:
@@ -468,10 +471,13 @@ def forward(
 
         if not torch.jit.is_scripting():
             # x from B x C x ... to C x B x ...
-            perm = np.arange(len(x.shape))
-            perm[0], perm[1] = perm[1], perm[0]
-            # Compute inverse permutation
-            inv_perm = np.arange(len(x.shape))[np.argsort(perm)]
+            perm = torch.cat(
+                (
+                    torch.tensor([1, 0], dtype=torch.long, device=x.device),
+                    torch.arange(2, x.ndim, dtype=torch.long, device=x.device),
+                )
+            )
+            inv_perm = perm
         else:
             raise NotImplementedError()
             # TorchScript in 2D for static inference

compressai/models/base.py

@@ -39,7 +39,7 @@
 
 from compressai.entropy_models import EntropyBottleneck, GaussianConditional
 from compressai.latent_codecs import LatentCodec
-from compressai.models.utils import update_registered_buffers
+from compressai.models.utils import remap_old_keys, update_registered_buffers
 
 __all__ = [
     "CompressionModel",
@@ -103,6 +103,7 @@ def load_state_dict(self, state_dict, strict=True):
                     ["_quantized_cdf", "_offset", "_cdf_length"],
                     state_dict,
                 )
+                state_dict = remap_old_keys(name, state_dict)
 
             if isinstance(module, GaussianConditional):
                 update_registered_buffers(

compressai/models/utils.py

@@ -27,10 +27,14 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+from collections import OrderedDict
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+KEY_MAP = {"_bias": "biases", "_matrix": "matrices", "_factor": "factors"}
+
 
 def find_named_module(module, query):
     """Helper function to find a named module. Returns a `nn.Module` or `None`
@@ -125,6 +129,24 @@ def update_registered_buffers(
         )
 
 
+def remap_old_keys(module_name, state_dict):
+    def remap_subkey(s: str) -> str:
+        for k, v in KEY_MAP.items():
+            if s.startswith(k):
+                return ".".join((v, s.split(k)[1]))
+
+        return s
+
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        if k.startswith(module_name):
+            k = ".".join((module_name, remap_subkey(k.split(f"{module_name}.")[1])))
+
+        new_state_dict[k] = v
+
+    return new_state_dict
+
+
 def conv(in_channels, out_channels, kernel_size=5, stride=2):
     return nn.Conv2d(
         in_channels,

tests/test_entropy_models.py

@@ -32,6 +32,8 @@
 import pytest
 import torch
 
+from packaging import version
+
 from compressai.entropy_models import (
     EntropyBottleneck,
     EntropyModel,
@@ -242,6 +244,33 @@ def test_loss(self):
     #     assert torch.allclose(y0[0], y1[0])
     #     assert torch.all(y1[1] == 0)  # not yet supported
 
+    @pytest.mark.skipif(
+        version.parse(torch.__version__) < version.parse("2.0.0"),
+        reason="torch.compile only available for torch>=2.0",
+    )
+    def test_compiling(self):
+        entropy_bottleneck = EntropyBottleneck(128)
+        x0 = torch.rand(1, 128, 32, 32)
+        x1 = x0.clone()
+        x0.requires_grad_(True)
+        x1.requires_grad_(True)
+
+        torch.manual_seed(32)
+        y0 = entropy_bottleneck(x0)
+
+        m = torch.compile(entropy_bottleneck)
+
+        torch.manual_seed(32)
+        y1 = m(x1)
+
+        assert torch.allclose(y0[0], y1[0])
+        assert torch.allclose(y0[1], y1[1])
+
+        y0[0].sum().backward()
+        y1[0].sum().backward()
+
+        assert torch.allclose(x0.grad, x1.grad)
+
     def test_update(self):
         # get a pretrained model
         net = bmshj2018_factorized(quality=1, pretrained=True).eval()