pytorch import

requirements.txt

@@ -5,4 +5,5 @@ onnxruntime==1.18.1
 tensorflow==2.12.0
 pystache==0.6.5
 traits~=7.0
-onnxscript==0.3.2
+onnxscript==0.3.2
+torch==2.9.1 --index-url https://download.pytorch.org/whl/cpu

src/acetone_nnet/generator/neural_network.py

@@ -69,7 +69,9 @@
     Softmax,
 )
 
-MODEL_TYPE = str | Path | onnx.ModelProto
+from torch.export import ExportedProgram
+
+MODEL_TYPE = str | Path | onnx.ModelProto | ExportedProgram
 if (3, 12) > version_info >= (3, 10):
     MODEL_TYPE = MODEL_TYPE | Sequential | Functional
 

src/acetone_nnet/importers/PYTORCH_importer/parse_pytorch.py

@@ -0,0 +1,45 @@
+from pathlib import Path
+
+import numpy as np
+
+from acetone_nnet.generator import (
+    ActivationFunctions,
+    Add,
+    Average,
+    AveragePooling2D,
+    BatchNormalization,
+    Concatenate,
+    ConstantPad,
+    Conv2D,
+    Dense,
+    Flatten,
+    InputLayer,
+    LeakyReLu,
+    Linear,
+    Maximum,
+    MaxPooling2D,
+    Minimum,
+    Multiply,
+    ReLu,
+    ResizeCubic,
+    ResizeLinear,
+    ResizeNearest,
+    Sigmoid,
+    Softmax,
+    Subtract,
+    TanH,
+)
+from acetone_nnet.graph import graph_interpretor
+from acetone_nnet.ir import Layer
+
+from torch.export import export, ExportedProgram
+def load_pytorch(program : ExportedProgram):
+    layers: list[Layer] = []
+    layers.append(InputLayer(
+        original_name="",
+        idx=0,
+        size=3,
+        input_shape=[1,1,1,3],
+        data_format="channels_first",
+    ))
+    return layers, "f4", "f4", "channels_first", 10, {}

src/acetone_nnet/importers/parser.py

@@ -31,6 +31,7 @@
 
 from acetone_nnet.ir import Layer
 
+from torch.export import ExportedProgram
 
 def parser(
         file_to_parse: Any,
@@ -78,7 +79,11 @@ def parser(
     if type(file_to_parse) is onnx.ModelProto:
         from .ONNX_importer.parser_ONNX import load_onnx
         return load_onnx(file_to_parse)
-
+
+    if type(file_to_parse) is ExportedProgram:
+        from .PYTORCH_importer.parse_pytorch import load_pytorch
+        return load_pytorch(file_to_parse)
+
     # TODO Conditional import/test if keras is available
     if ((3, 12) > version_info >= (3, 10) and
             (type(file_to_parse) is Functional or type(file_to_parse) is Sequential)):

tests/tests_importer/importerTestCase.py

@@ -27,7 +27,7 @@
 
 import acetone_nnet
 from acetone_nnet import Layer
-
+from torch.export import ExportedProgram
 
 def create_initializer_tensor(
     name: str,
@@ -57,7 +57,7 @@ def tearDown(self) -> None:
 
     def import_layers(
         self,
-        file: str | onnx.ModelProto | Functional | Sequential,
+        file: str | onnx.ModelProto | Functional | Sequential | ExportedProgram,
         versions: dict[int, str] | dict[str, str] | None = None,
     ) -> acetone_nnet.generator.CodeGenerator:
         """Create the CodeGenerator object from the file."""

tests/tests_importer/tests_pytorch/__init__.py

@@ -0,0 +1,18 @@
+"""*******************************************************************************
+* ACETONE: Predictable programming framework for ML applications in safety-critical systems
+* Copyright (c) 2022. ONERA
+* This file is part of ACETONE
+*
+* ACETONE is free software ;
+* you can redistribute it and/or modify it under the terms of the GNU Lesser General Public
+* License as published by the Free Software Foundation ;
+* either version 3 of  the License, or (at your option) any later version.
+*
+* ACETONE is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY ;
+* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+* See the GNU Lesser General Public License for more details.
+*
+* You should have received a copy of the GNU Lesser General Public License along with this program ;
+* if not, write to the Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307  USA
+******************************************************************************
+"""

tests/tests_importer/tests_pytorch/test_Linear.py

@@ -0,0 +1,30 @@
+"""*******************************************************************************
+* ACETONE: Predictable programming framework for ML applications in safety-critical systems
+* Copyright (c) 2022. ONERA
+* This file is part of ACETONE
+*
+* ACETONE is free software ;
+* you can redistribute it and/or modify it under the terms of the GNU Lesser General Public
+* License as published by the Free Software Foundation ;
+* either version 3 of  the License, or (at your option) any later version.
+*
+* ACETONE is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY ;
+* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+* See the GNU Lesser General Public License for more details.
+*
+* You should have received a copy of the GNU Lesser General Public License along with this program ;
+* if not, write to the Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307  USA
+******************************************************************************
+"""
+
+from tests.tests_importer import importerTestCase
+from torch.nn import Linear, Sequential
+from torch.export import export, ExportedProgram
+import torch
+class TestLinear(importerTestCase.ImporterTestCase):
+    """Test for MatMul Layer."""
+
+    def test_linear_1(self) -> None:
+        model = Sequential(*[Linear(3,4)])
+        program:ExportedProgram = export(model,(torch.rand(3),))
+        cgen = self.import_layers(program)