chores: update tested pytorch images and bump dependencies

(cherry picked from commit 3c8b7bc01f1401c8e0a08934a317e57d7299365b)

Author: fracape

.gitlab-ci.yml

@@ -96,7 +96,7 @@ test:
   parallel:
     matrix:
       - PYTORCH_IMAGE:
-          - "1.8.1-cuda11.1-cudnn8-devel"
+          - "1.13.1-cuda11.6-cudnn8-devel"
           - "2.2.0-cuda11.8-cudnn8-devel"
   tags:
     - docker

compressai/models/video/google.py

@@ -35,7 +35,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from torch.cuda import amp
+from torch import amp
 
 from compressai.entropy_models import EntropyBottleneck, GaussianConditional
 from compressai.layers import QReLU
@@ -350,7 +350,6 @@ def gaussian_volume(x, sigma: float, num_levels: int):
             volume.append(interp.unsqueeze(2))
         return torch.cat(volume, dim=2)
 
-    @amp.autocast(enabled=False)
     def warp_volume(self, volume, flow, scale_field, padding_mode: str = "border"):
         """3D volume warping."""
         if volume.ndimension() != 5:
@@ -360,14 +359,18 @@ def warp_volume(self, volume, flow, scale_field, padding_mode: str = "border"):
 
         N, C, _, H, W = volume.size()
 
-        grid = meshgrid2d(N, C, H, W, volume.device)
-        update_grid = grid + flow.permute(0, 2, 3, 1).float()
-        update_scale = scale_field.permute(0, 2, 3, 1).float()
-        volume_grid = torch.cat((update_grid, update_scale), dim=-1).unsqueeze(1)
-
-        out = F.grid_sample(
-            volume.float(), volume_grid, padding_mode=padding_mode, align_corners=False
-        )
+        with amp.autocast(device_type=volume.device.type, enabled=False):
+            grid = meshgrid2d(N, C, H, W, volume.device)
+            update_grid = grid + flow.permute(0, 2, 3, 1).float()
+            update_scale = scale_field.permute(0, 2, 3, 1).float()
+            volume_grid = torch.cat((update_grid, update_scale), dim=-1).unsqueeze(1)
+
+            out = F.grid_sample(
+                volume.float(),
+                volume_grid,
+                padding_mode=padding_mode,
+                align_corners=False,
+            )
         return out.squeeze(2)
 
     def forward_prediction(self, x_ref, motion_info):

compressai/utils/video/eval_model/__main__.py

@@ -43,8 +43,7 @@
 import torch.nn.functional as F
 
 from pytorch_msssim import ms_ssim
-from torch import Tensor
-from torch.cuda import amp
+from torch import Tensor, amp
 from torch.utils.model_zoo import tqdm
 
 import compressai
@@ -370,7 +369,9 @@ def run_inference(
         if sequence_metrics_path.is_file():
             continue
 
-        with amp.autocast(enabled=args["half"]):
+        with amp.autocast(
+            device_type=next(net.parameters()).device.type, enabled=args["half"]
+        ):
             with torch.no_grad():
                 if entropy_estimation:
                     metrics = eval_model_entropy_estimation(net, filepath)

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=42", "wheel", "pybind11>=2.6.0", "torch"]
+requires = ["setuptools>=42", "wheel", "pybind11>=2.12", "torch"]
 build-backend = "setuptools.build_meta"
 
 [tool.black]

setup.py

@@ -136,11 +136,11 @@ def get_extra_requirements():
     python_requires=">=3.8",
     install_requires=[
         "einops",
-        "numpy>=1.21.0",
+        "numpy>=1.21.0, <2",
         "pandas",
         "scipy",
         "matplotlib",
-        "torch>=1.7.1, <2.3",
+        "torch>=1.13.1",
         "torch-geometric>=2.3.0",
         "typing-extensions>=4.0.0",
         "torchvision",

tests/test_entropy_models.py

@@ -273,6 +273,7 @@ def test_compiling(self):
 
     def test_update(self):
         # get a pretrained model
+
         net = bmshj2018_factorized(quality=1, pretrained=True).eval()
         assert not net.update()
         assert not net.update(force=False)