Setup: CI tests against pytorch 2.x

.github/workflows/develop_install.yml

@@ -16,7 +16,7 @@ jobs:
 
       matrix:
         python_version: ['3.8', '3.9']
-        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0', '2.0.1', '2.1.0']
         platform: ['windows-latest', 'ubuntu-latest', 'macos-latest']
 
 

.github/workflows/end_to_end.yml

@@ -16,7 +16,7 @@ jobs:
 
       matrix:
         python_version: ['3.8', '3.9']
-        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0', '2.0.1', '2.1.0']
         platform: ['windows-latest', 'ubuntu-latest', 'macos-latest']
 
 

.github/workflows/examples_pytest.yml

@@ -16,7 +16,7 @@ jobs:
 
       matrix:
         python_version: ['3.8', '3.9']
-        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0', '2.0.1', '2.1.0']
         platform: ['windows-latest', 'ubuntu-latest', 'macos-latest']
         jit_status: ['jit_disabled', 'jit_enabled']
 

.github/workflows/finn_integration.yml

@@ -16,7 +16,7 @@ jobs:
 
       matrix:
         python_version: ['3.8', '3.9']
-        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0', '2.0.1', '2.1.0']
         platform: ['windows-latest', 'ubuntu-latest']
 
 

.github/workflows/gen_github_actions.py

@@ -17,7 +17,7 @@
 # Reduced Test for PRs, except when a review is requested
 PYTHON_VERSIONS_REDUCED = ('3.8',)
 
-PYTORCH_LIST_REDUCED = ["1.9.1", "1.10.1", "1.13.0"]
+PYTORCH_LIST_REDUCED = ["1.9.1", "1.13.0", "2.1.0"]
 
 PLATFORM_LIST_REDUCED = ['ubuntu-latest']
 
@@ -35,7 +35,7 @@
 # Data shared betwen Nox sessions and Github Actions, formatted as tuples
 PYTHON_VERSIONS = ('3.8', '3.9')
 
-PYTORCH_VERSIONS = ('1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0')
+PYTORCH_VERSIONS = ('1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0', '2.0.1', '2.1.0')
 JIT_STATUSES = ('jit_disabled', 'jit_enabled')
 
 # Data used only by Github Actions, formatted as lists or lists of ordered dicts

.github/workflows/notebook.yml

@@ -16,7 +16,7 @@ jobs:
 
       matrix:
         python_version: ['3.8', '3.9']
-        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0', '2.0.1', '2.1.0']
         platform: ['windows-latest', 'ubuntu-latest', 'macos-latest']
 
 

.github/workflows/ort_integration.yml

@@ -16,7 +16,7 @@ jobs:
 
       matrix:
         python_version: ['3.8', '3.9']
-        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0', '2.0.1', '2.1.0']
         platform: ['windows-latest', 'ubuntu-latest', 'macos-latest']
 
 

.github/workflows/pytest.yml

@@ -16,7 +16,7 @@ jobs:
 
       matrix:
         python_version: ['3.8', '3.9']
-        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.10.1', '1.11.0', '1.12.1', '1.13.0', '2.0.1', '2.1.0']
         platform: ['windows-latest', 'ubuntu-latest', 'macos-latest']
         jit_status: ['jit_disabled', 'jit_enabled']
 

.github/workflows/reduced_develop_install.yml

@@ -18,7 +18,7 @@ jobs:
 
       matrix:
         python_version: ['3.8']
-        pytorch_version: ['1.9.1', '1.10.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.13.0', '2.1.0']
         platform: ['ubuntu-latest']
 
 

.github/workflows/reduced_end_to_end.yml

@@ -18,7 +18,7 @@ jobs:
 
       matrix:
         python_version: ['3.8']
-        pytorch_version: ['1.9.1', '1.10.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.13.0', '2.1.0']
         platform: ['ubuntu-latest']
 
 

.github/workflows/reduced_examples_pytest.yml

@@ -18,7 +18,7 @@ jobs:
 
       matrix:
         python_version: ['3.8']
-        pytorch_version: ['1.9.1', '1.10.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.13.0', '2.1.0']
         platform: ['ubuntu-latest']
         jit_status: ['jit_disabled']
 

.github/workflows/reduced_finn_integration.yml

@@ -18,7 +18,7 @@ jobs:
 
       matrix:
         python_version: ['3.8']
-        pytorch_version: ['1.9.1', '1.10.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.13.0', '2.1.0']
         platform: ['ubuntu-latest']
 
 

.github/workflows/reduced_notebook.yml

@@ -18,7 +18,7 @@ jobs:
 
       matrix:
         python_version: ['3.8']
-        pytorch_version: ['1.9.1', '1.10.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.13.0', '2.1.0']
         platform: ['ubuntu-latest']
 
 

.github/workflows/reduced_ort_integration.yml

@@ -18,7 +18,7 @@ jobs:
 
       matrix:
         python_version: ['3.8']
-        pytorch_version: ['1.9.1', '1.10.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.13.0', '2.1.0']
         platform: ['ubuntu-latest']
 
 

.github/workflows/reduced_pytest.yml

@@ -18,7 +18,7 @@ jobs:
 
       matrix:
         python_version: ['3.8']
-        pytorch_version: ['1.9.1', '1.10.1', '1.13.0']
+        pytorch_version: ['1.9.1', '1.13.0', '2.1.0']
         platform: ['ubuntu-latest']
         jit_status: ['jit_disabled']
 

noxfile.py

@@ -24,7 +24,9 @@
     '1.10.1': '0.11.2',
     '1.11.0': '0.12.0',
     '1.12.1': '0.13.1',
-    '1.13.0': '0.14.0'}
+    '1.13.0': '0.14.0',
+    '2.0.1': '0.15.2',
+    '2.1.0': '0.16.0'}
 
 PARSED_TORCHVISION_VERSION_DICT = {version.parse(k): v for k, v in TORCHVISION_VERSION_DICT.items()}
 

src/brevitas/export/onnx/debug.py

@@ -14,6 +14,7 @@ class DebugMarkerFunction(Function):
     @staticmethod
     def symbolic(g, input, export_debug_name):
         ret = g.op('brevitas.onnx::DebugMarker', input, export_debug_name_s=export_debug_name)
+        ret.setType(input.type())
         return ret
 
     @staticmethod

src/brevitas/export/onnx/qonnx/function.py

@@ -3,6 +3,7 @@
 
 import torch
 from torch.autograd import Function
+from torch.onnx.symbolic_helper import _get_tensor_sizes
 
 from brevitas.core.bit_width import BitWidthConst
 from brevitas.core.function_wrapper.clamp import TensorClamp
@@ -19,6 +20,7 @@ class BrevitasBinaryQuantFn(Function):
     @staticmethod
     def symbolic(g, x, scale, zero_point, bit_width, narrow_range, signed, rounding_mode):
         ret = g.op(f'{DOMAIN_STRING}::BipolarQuant', x, scale)
+        ret.setType(x.type())
         return ret
 
     @staticmethod
@@ -40,6 +42,7 @@ def symbolic(g, x, scale, zero_point, bit_width, narrow_range, signed, rounding_
             rounding_mode_s=rounding_mode,
             signed_i=int(signed),
             narrow_i=int(narrow_range))
+        ret.setType(x.type())
         return ret
 
     @staticmethod
@@ -66,6 +69,7 @@ def symbolic(g, x, scale, zero_point, input_bit_width, output_bit_width, roundin
             input_bit_width,
             output_bit_width,
             rounding_mode_s=rounding_mode)
+        ret.setType(x.type())
         return ret
 
     @staticmethod

src/brevitas/nn/quant_activation.py

@@ -13,6 +13,24 @@
 from .quant_layer import QuantNonLinearActLayer as QuantNLAL
 
 
+# Starting from Torch 2.0, nn.Module init function accepts custom *args and **kwargs
+# torch.nn.Sigmoid does not provide its own init method, and the presence of *args + **kwargs
+# conflicts with the dependency injection package
+class Sigmoid(nn.Sigmoid):
+
+    def __init__(self):
+        super().__init__()
+
+
+# Starting from Torch 2.0, nn.Module init function accepts custom *args and **kwargs
+# torch.nn.Tanh does not provide its own init method, and the presence of *args + **kwargs
+# conflicts with the dependency injection package
+class Tanh(nn.Tanh):
+
+    def __init__(self):
+        super().__init__()
+
+
 class QuantReLU(QuantNLAL):
 
     def __init__(
@@ -41,7 +59,7 @@ def __init__(
             **kwargs):
         QuantNLAL.__init__(
             self,
-            act_impl=nn.Sigmoid,
+            act_impl=Sigmoid,
             passthrough_act=False,
             input_quant=input_quant,
             act_quant=act_quant,
@@ -59,7 +77,7 @@ def __init__(
             **kwargs):
         QuantNLAL.__init__(
             self,
-            act_impl=nn.Tanh,
+            act_impl=Tanh,
             passthrough_act=False,
             input_quant=input_quant,
             act_quant=act_quant,

tests/brevitas/graph/equalization_fixtures.py

@@ -109,6 +109,10 @@ def linearmha_model(bias, add_bias_kv, batch_first):
     if torch_version < version.parse('1.9.1'):
         pytest.skip(f"batch_first not supported in MHA with torch version {torch_version}")
 
+    # Skip due to following issue https://github.com/pytorch/pytorch/issues/97128
+    if torch_version == version.parse('2.0.1') and not bias and batch_first and not add_bias_kv:
+        pytest.skip(f"Skip due to a regression in pytorch 2.0.1")
+
     class LinearMhaModel(nn.Module):
 
         def __init__(self) -> None:
@@ -135,6 +139,10 @@ def layernormmha_model(bias, add_bias_kv, batch_first):
     if torch_version < version.parse('1.9.1'):
         pytest.skip(f"batch_first not supported in MHA with torch version {torch_version}")
 
+    # Skip due to following issue https://github.com/pytorch/pytorch/issues/97128
+    if torch_version == version.parse('2.0.1') and not bias and batch_first and not add_bias_kv:
+        pytest.skip(f"Skip due to a regression in pytorch 2.0.1")
+
     class LayerNormMhaModel(nn.Module):
 
         def __init__(self) -> None:
@@ -164,6 +172,10 @@ def mhalinear_model(bias, add_bias_kv, batch_first):
     if torch_version < version.parse('1.9.1'):
         pytest.skip(f"batch_first not supported in MHA with torch version {torch_version}")
 
+    # Skip due to following issue https://github.com/pytorch/pytorch/issues/97128
+    if torch_version == version.parse('2.0.1') and not bias and batch_first and not add_bias_kv:
+        pytest.skip(f"Skip due to a regression in pytorch 2.0.1")
+
     class MhaLinearModel(nn.Module):
 
         def __init__(self) -> None:

tests/brevitas/graph/test_equalization.py

@@ -124,7 +124,9 @@ def test_models(toy_model, merge_bias, request):
     inp = torch.randn(in_shape)
 
     model.eval()
-    expected_out = model(inp)
+    with torch.no_grad():
+        expected_out = model(inp)
+
     model = symbolic_trace(model)
     regions = _extract_regions(model)
     scale_factor_regions = equalize_test(
@@ -135,7 +137,8 @@ def test_models(toy_model, merge_bias, request):
         scale_computation_type='maxabs')
     shape_scale_regions = [scale.shape for scale in scale_factor_regions]
 
-    out = model(inp)
+    with torch.no_grad():
+        out = model(inp)
     assert len(regions) > 0
     assert torch.allclose(expected_out, out, atol=ATOL)
     # Check that at least one region performs "true" equalization

tests/brevitas_end_to_end/test_torchvision_models.py

@@ -73,6 +73,11 @@ def torchvision_model(model_name, quantize_fn):
     if torch_version < version.parse('1.11.0') and model_name == 'vit_b_32':
         return None
 
+    # Due to a regression in torchvision, we cannot load pretrained weights for effnet_b0
+    # https://github.com/pytorch/vision/issues/7744
+    if torch_version == version.parse('2.1.0') and model_name == 'efficientnet_b0':
+        return None
+
     # Deeplab and fcn are in a different module, and they have a dict as output which is not suited for torchscript
     if model_name in ('deeplabv3_resnet50', 'fcn_resnet50'):
         model_fn = getattr(modelzoo.segmentation, model_name)