diff --git a/.github/workflows/develop_install.yml b/.github/workflows/develop_install.yml index c2bb054c3..eff4e15b0 100644 --- a/.github/workflows/develop_install.yml +++ b/.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'] diff --git a/.github/workflows/end_to_end.yml b/.github/workflows/end_to_end.yml index 8ae512f91..98d1a309e 100644 --- a/.github/workflows/end_to_end.yml +++ b/.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'] diff --git a/.github/workflows/examples_pytest.yml b/.github/workflows/examples_pytest.yml index 8b50b19bb..31ccb3ffa 100644 --- a/.github/workflows/examples_pytest.yml +++ b/.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'] diff --git a/.github/workflows/finn_integration.yml b/.github/workflows/finn_integration.yml index 8590e4b81..19becbe70 100644 --- a/.github/workflows/finn_integration.yml +++ b/.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'] diff --git a/.github/workflows/gen_github_actions.py b/.github/workflows/gen_github_actions.py index 535a777a4..0f996dc2b 100644 --- a/.github/workflows/gen_github_actions.py +++ b/.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 diff --git a/.github/workflows/notebook.yml b/.github/workflows/notebook.yml index ef5a20e41..fd079b0ca 100644 --- a/.github/workflows/notebook.yml +++ b/.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'] diff --git a/.github/workflows/ort_integration.yml b/.github/workflows/ort_integration.yml index e09432973..7d0eccb27 100644 --- a/.github/workflows/ort_integration.yml +++ b/.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'] diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 67b9eaccb..3d9f43ea9 100644 --- a/.github/workflows/pytest.yml +++ b/.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'] diff --git a/.github/workflows/reduced_develop_install.yml b/.github/workflows/reduced_develop_install.yml index e2640ab17..8d2c1b45a 100644 --- a/.github/workflows/reduced_develop_install.yml +++ b/.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'] diff --git a/.github/workflows/reduced_end_to_end.yml b/.github/workflows/reduced_end_to_end.yml index af3ef3a4d..d19469ef6 100644 --- a/.github/workflows/reduced_end_to_end.yml +++ b/.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'] diff --git a/.github/workflows/reduced_examples_pytest.yml b/.github/workflows/reduced_examples_pytest.yml index 15441ffd0..c10aa1c74 100644 --- a/.github/workflows/reduced_examples_pytest.yml +++ b/.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'] diff --git a/.github/workflows/reduced_finn_integration.yml b/.github/workflows/reduced_finn_integration.yml index eda43d79f..c7ebafe8c 100644 --- a/.github/workflows/reduced_finn_integration.yml +++ b/.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'] diff --git a/.github/workflows/reduced_notebook.yml b/.github/workflows/reduced_notebook.yml index f53ba8244..a7293e4d3 100644 --- a/.github/workflows/reduced_notebook.yml +++ b/.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'] diff --git a/.github/workflows/reduced_ort_integration.yml b/.github/workflows/reduced_ort_integration.yml index e0eb0a655..1671d6a13 100644 --- a/.github/workflows/reduced_ort_integration.yml +++ b/.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'] diff --git a/.github/workflows/reduced_pytest.yml b/.github/workflows/reduced_pytest.yml index 19b1d31f2..12e3d360c 100644 --- a/.github/workflows/reduced_pytest.yml +++ b/.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'] diff --git a/noxfile.py b/noxfile.py index 8c8f56284..ffb1c5fbd 100644 --- a/noxfile.py +++ b/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()} diff --git a/src/brevitas/export/onnx/debug.py b/src/brevitas/export/onnx/debug.py index 85acba31b..3d5c236ac 100644 --- a/src/brevitas/export/onnx/debug.py +++ b/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 diff --git a/src/brevitas/export/onnx/qonnx/function.py b/src/brevitas/export/onnx/qonnx/function.py index cd1642d11..d410ee31e 100644 --- a/src/brevitas/export/onnx/qonnx/function.py +++ b/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 diff --git a/src/brevitas/nn/quant_activation.py b/src/brevitas/nn/quant_activation.py index 7c81fbf8a..c8daa59ef 100644 --- a/src/brevitas/nn/quant_activation.py +++ b/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, diff --git a/tests/brevitas/graph/equalization_fixtures.py b/tests/brevitas/graph/equalization_fixtures.py index 128f3953a..8e13c49cd 100644 --- a/tests/brevitas/graph/equalization_fixtures.py +++ b/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: diff --git a/tests/brevitas/graph/test_equalization.py b/tests/brevitas/graph/test_equalization.py index 7ae9022b7..caca0fd29 100644 --- a/tests/brevitas/graph/test_equalization.py +++ b/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 diff --git a/tests/brevitas_end_to_end/test_torchvision_models.py b/tests/brevitas_end_to_end/test_torchvision_models.py index b554c8113..e8c34d961 100644 --- a/tests/brevitas_end_to_end/test_torchvision_models.py +++ b/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)