Skip to content

Multiple Inputs Support #1431

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yarden-yagil-sony merged 19 commits into from
May 7, 2025
Merged

Multiple Inputs Support #1431

Changes from 1 commit
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
e343846
dynamic batch axis in input nodes
Apr 29, 2025
2053127
minor fix
May 5, 2025
4a35021
add multiple inputs test
May 5, 2025
153b9a4
move export tests to integration tests
May 5, 2025
f371c15
change input names
May 6, 2025
c64fe07
update onnx opset version
May 6, 2025
80b0105
conditional onnx version
May 6, 2025
254da4c
conditional onnx version
May 6, 2025
596068a
cleanup
May 6, 2025
44f5896
update torch test 2.2 to 2.6
May 6, 2025
c71f3ea
Merge branch 'refs/heads/main' into multiple_inputs_support
May 6, 2025
02e2839
update test
May 6, 2025
91015ef
fix test to work with torch 2.6
May 6, 2025
1e4de4e
update test torch version to 2.6
May 6, 2025
dda57fb
update test torch version to 2.6
May 6, 2025
097dfe4
torch.load weights only due to torch version 2.6
May 7, 2025
9d10849
update readme
May 7, 2025
538ace7
Merge branch 'refs/heads/main' into multiple_inputs_support
May 7, 2025
74b7ac7
update readme torch table
May 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions ...compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -118,14 +120,17 @@ def export(self, output_names=None) -> None:
dynamic_axes = {'input': {0: 'batch_size'}}
dynamic_axes.update({name: {0: 'batch_size'} for name in output_names})

input_names = [n.name for n in self.model.node_sort if n.type == DummyPlaceHolder]
dynamic_axes.update({name: {0: 'batch_size'} for name in input_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())
Expand All @@ -137,7 +142,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)

Expand Down