Multiple Inputs Support | Onnx Opset V20 | Add torch2.6 & Remove torch2.2 (#1431)

Author: yarden-yagil-sony

…kflows/run_tests_python310_pytorch22.yml …kflows/run_tests_python310_pytorch26.yml.github/workflows/run_tests_python310_pytorch22.yml renamed to .github/workflows/run_tests_python310_pytorch26.yml .github/workflows/run_tests_python310_pytorch22.yml renamed to .github/workflows/run_tests_python310_pytorch26.yml

@@ -1,4 +1,4 @@
-name: Python 3.10, Pytorch 2.2
+name: Python 3.10, Pytorch 2.6
 on:
   workflow_dispatch: # Allow manual triggers
   schedule:
@@ -16,4 +16,4 @@ jobs:
     uses: ./.github/workflows/run_pytorch_tests.yml
     with:
       python-version: "3.10"
-      torch-version: "2.2.*"
+      torch-version: "2.6.*"

…kflows/run_tests_python311_pytorch22.yml …kflows/run_tests_python311_pytorch26.yml.github/workflows/run_tests_python311_pytorch22.yml renamed to .github/workflows/run_tests_python311_pytorch26.yml .github/workflows/run_tests_python311_pytorch22.yml renamed to .github/workflows/run_tests_python311_pytorch26.yml

@@ -1,4 +1,4 @@
-name: Python 3.11, Pytorch 2.2
+name: Python 3.11, Pytorch 2.6
 on:
   workflow_dispatch: # Allow manual triggers
   schedule:
@@ -16,4 +16,4 @@ jobs:
     uses: ./.github/workflows/run_pytorch_tests.yml
     with:
       python-version: "3.11"
-      torch-version: "2.2.*"
+      torch-version: "2.6.*"

…kflows/run_tests_python312_pytorch22.yml …kflows/run_tests_python312_pytorch26.yml.github/workflows/run_tests_python312_pytorch22.yml renamed to .github/workflows/run_tests_python312_pytorch26.yml .github/workflows/run_tests_python312_pytorch22.yml renamed to .github/workflows/run_tests_python312_pytorch26.yml

@@ -1,4 +1,4 @@
-name: Python 3.12, Pytorch 2.2
+name: Python 3.12, Pytorch 2.6
 on:
   workflow_dispatch: # Allow manual triggers
   schedule:
@@ -16,4 +16,4 @@ jobs:
     uses: ./.github/workflows/run_pytorch_tests.yml
     with:
       python-version: "3.12"
-      torch-version: "2.2.*"
+      torch-version: "2.6.*"

…rkflows/run_tests_python39_pytorch22.yml …rkflows/run_tests_python39_pytorch26.yml.github/workflows/run_tests_python39_pytorch22.yml renamed to .github/workflows/run_tests_python39_pytorch26.yml .github/workflows/run_tests_python39_pytorch22.yml renamed to .github/workflows/run_tests_python39_pytorch26.yml

@@ -1,4 +1,4 @@
-name: Python 3.9, Pytorch 2.2
+name: Python 3.9, Pytorch 2.6
 on:
   workflow_dispatch: # Allow manual triggers
   schedule:
@@ -16,4 +16,4 @@ jobs:
     uses: ./.github/workflows/run_pytorch_tests.yml
     with:
       python-version: "3.9"
-      torch-version: "2.2.*"
+      torch-version: "2.6.*"

README.md

@@ -17,7 +17,7 @@ ______________________________________________________________________
   <a href="#license">License</a>
 </p>
 <p align="center">
-  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/pytorch-2.2%20%7C%202.3%20%7C%202.4%20%7C%202.5-blue" /></a>
+  <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/pytorch-2.3%20%7C%202.4%20%7C%202.5%20%7C%202.6-blue" /></a>
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/tensorflow-2.14%20%7C%202.15-blue" /></a>
   <a href="https://sony.github.io/model_optimization#prerequisites"><img src="https://img.shields.io/badge/python-3.9%20%7C%203.10%20%7C%203.11%20%7C%203.12-blue" /></a>
   <a href="https://github.com/sony/model_optimization/releases"><img src="https://img.shields.io/github/v/release/sony/model_optimization" /></a>
@@ -30,7 +30,7 @@ ________________________________________________________________________________
 
 ## <div align="center">Getting Started</div>
 ### Quick Installation
-Pip install the model compression toolkit package in a Python>=3.9 environment with PyTorch>=2.1 or Tensorflow>=2.14.
+Pip install the model compression toolkit package in a Python>=3.9 environment with PyTorch>=2.3 or Tensorflow>=2.14.
 ```
 pip install model-compression-toolkit
 ```
@@ -130,12 +130,12 @@ Currently, MCT is being tested on various Python, Pytorch and TensorFlow version
 <details id="supported-versions">
   <summary>Supported Versions Table</summary>
 
-|                                                                                                                                                                                                                    | PyTorch 2.2                                                                                                                                                                                                              | PyTorch 2.3                                                                                                                                                                                                              | PyTorch 2.4                                                                                                                                                                                                              | PyTorch 2.5                                                                                                                                                                                                              |
+|                                                                                                                                                                                                                   | PyTorch 2.3                                                                                                                                                                                                              | PyTorch 2.4                                                                                                                                                                                                              | PyTorch 2.5                                                                                                                                                                                                              | PyTorch 2.6                                                                                                                                                                                                              |
 |-------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| Python 3.9   | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch22.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch22.yml)   | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch23.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch23.yml)   | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch24.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch24.yml)   | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch25.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch25.yml)   |
-| Python 3.10 | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch22.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch22.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch23.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch23.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch24.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch24.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch25.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch25.yml) |
-| Python 3.11 | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch22.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch22.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch23.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch23.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch24.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch24.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch25.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch25.yml) |
-| Python 3.12 | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch22.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch22.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch23.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch23.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch24.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch24.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch25.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch25.yml) |
+| Python 3.9  | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch23.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch23.yml)   | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch24.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch24.yml)   | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch25.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch25.yml)   | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch26.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python39_pytorch26.yml)   |
+| Python 3.10 | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch23.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch23.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch24.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch24.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch25.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch25.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch26.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python310_pytorch26.yml) |
+| Python 3.11 | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch23.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch23.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch24.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch24.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch25.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch25.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch26.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python311_pytorch26.yml) |
+| Python 3.12 | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch23.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch23.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch24.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch24.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch25.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch25.yml) | [![Run Tests](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch26.yml/badge.svg)](https://github.com/sony/model_optimization/actions/workflows/run_tests_python312_pytorch26.yml) |
 
 |             | TensorFlow 2.14                                                                                                                                                                                                        | TensorFlow 2.15                                                                                                                                                                                                        |
 |-------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -18,6 +18,8 @@
 import torch.nn
 
 from mct_quantizers import PytorchActivationQuantizationHolder, PytorchQuantizationWrapper
+
+from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlaceHolder
 from model_compression_toolkit.verify_packages import FOUND_ONNX
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
@@ -98,15 +100,17 @@ def export(self, output_names=None) -> None:
             model_output = self.model(*model_input) if isinstance(model_input, (list, tuple)) else self.model(
                 model_input)
 
+            input_nodes = [n for n in self.model.node_sort if n.type == DummyPlaceHolder]
+            input_names = [f"input_{i}" for i in range(len(input_nodes))] if len(input_nodes) > 1 else ["input"]
+            dynamic_axes = {name: {0: 'batch_size'} for name in input_names}
             if output_names is None:
                 # Determine number of outputs and prepare output_names and dynamic_axes
                 if isinstance(model_output, (list, tuple)):
                     output_names = [f"output_{i}" for i in range(len(model_output))]
-                    dynamic_axes = {'input': {0: 'batch_size'}}
                     dynamic_axes.update({name: {0: 'batch_size'} for name in output_names})
                 else:
                     output_names = ['output']
-                    dynamic_axes = {'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
+                    dynamic_axes.update({'output': {0: 'batch_size'}})
             else:
                 if isinstance(model_output, (list, tuple)):
                     num_of_outputs = len(model_output)
@@ -115,17 +119,15 @@ def export(self, output_names=None) -> None:
                 assert len(output_names) == num_of_outputs, (f"Mismatch between number of requested output names "
                                                              f"({output_names}) and model output count "
                                                              f"({num_of_outputs}):\n")
-                dynamic_axes = {'input': {0: 'batch_size'}}
                 dynamic_axes.update({name: {0: 'batch_size'} for name in output_names})
-
             if hasattr(self.model, 'metadata'):
                 onnx_bytes = BytesIO()
                 torch.onnx.export(self.model,
                                   tuple(model_input) if isinstance(model_input, list) else model_input,
                                   onnx_bytes,
                                   opset_version=self._onnx_opset_version,
                                   verbose=False,
-                                  input_names=['input'],
+                                  input_names=input_names,
                                   output_names=output_names,
                                   dynamic_axes=dynamic_axes)
                 onnx_model = onnx.load_from_string(onnx_bytes.getvalue())
@@ -137,7 +139,7 @@ def export(self, output_names=None) -> None:
                                   self.save_model_path,
                                   opset_version=self._onnx_opset_version,
                                   verbose=False,
-                                  input_names=['input'],
+                                  input_names=input_names,
                                   output_names=output_names,
                                   dynamic_axes=dynamic_axes)
 

model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 from typing import Callable
+from packaging import version
 
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.exporter.model_exporter.fw_agonstic.quantization_format import QuantizationFormat
@@ -30,6 +31,9 @@
     from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_torchscript_pytorch_exporter import FakelyQuantTorchScriptPyTorchExporter
     from model_compression_toolkit.exporter.model_wrapper.pytorch.validate_layer import is_pytorch_layer_exportable
 
+    if version.parse(torch.__version__) >= version.parse("2.4"):
+        DEFAULT_ONNX_OPSET_VERSION = 20
+
     supported_serialization_quantization_export_dict = {
         PytorchExportSerializationFormat.TORCHSCRIPT: [QuantizationFormat.FAKELY_QUANT],
         PytorchExportSerializationFormat.ONNX: [QuantizationFormat.FAKELY_QUANT, QuantizationFormat.MCTQ]

tests/pytorch_tests/exporter_tests/test_exporting_qat_models.py

@@ -51,7 +51,7 @@ def get_tmp_filepath(self):
         return tempfile.mkstemp('.pt')[1]
 
     def load_exported_model(self, filepath):
-        return torch.load(filepath)
+        return torch.load(filepath, weights_only=False)
 
     def infer(self, model, images):
         return model(images)
@@ -74,15 +74,15 @@ def export_qat_model(self):
         # Assert qat_ready can be saved and loaded
         _qat_ready_model_tmp_filepath = tempfile.mkstemp('.pt')[1]
         torch.save(self.qat_ready, _qat_ready_model_tmp_filepath)
-        self.qat_ready = torch.load(_qat_ready_model_tmp_filepath)
+        self.qat_ready = torch.load(_qat_ready_model_tmp_filepath, weights_only=False)
 
         self.final_model = mct.qat.pytorch_quantization_aware_training_finalize_experimental(self.qat_ready)
 
         # Assert final_model can be used for inference, can be saved and loaded
         self.final_model(images[0])
         _final_model_tmp_filepath = tempfile.mkstemp('.pt')[1]
         torch.save(self.final_model, _final_model_tmp_filepath)
-        self.final_model = torch.load(_final_model_tmp_filepath)
+        self.final_model = torch.load(_final_model_tmp_filepath, weights_only=False)
         self.final_model(images[0])
 
         self.filepath = self.get_tmp_filepath()

tests/pytorch_tests/function_tests/test_fully_quantized_exporter.py

@@ -94,7 +94,7 @@ def test_save_and_load_model(self):
         print(f"Float Pytorch .pth Model: {float_model_file}")
 
         model = copy.deepcopy(self.fully_quantized_mbv2)
-        model.load_state_dict(torch.load(model_file))
+        model.load_state_dict(torch.load(model_file, weights_only=False))
         model.eval()
         model(next(self.representative_data_gen()))