Python pnnx with pnnx binary (#5067)

.gitignore

@@ -54,6 +54,7 @@ __pycache__
 *.pyd
 *.egg-info/
 python/setup.py
+tools/pnnx/python/setup.py
 
 # Clangd
 .cache/

tools/pnnx/python/README.md

@@ -0,0 +1,199 @@
+# pnnx
+python wrapper of pnnx, only support python 3.7+  now.
+
+Install from pip
+==================
+
+pnnx is available as wheel packages for macOS, Windows and Linux distributions, you can install with pip:
+
+```
+pip install pnnx
+```
+
+# Build & Install from source
+
+## Prerequisites
+
+**On Unix (Linux, OS X)**
+
+* A compiler with C++14 support
+* CMake >= 3.4
+
+**On Mac**
+
+* A compiler with C++14 support
+* CMake >= 3.4
+
+**On Windows**
+
+* Visual Studio 2015 or higher
+* CMake >= 3.4
+
+## Build & install
+1. clone ncnn.
+```bash
+git clone https://github.com/Tencent/ncnn.git
+```
+2. install pytorch 
+
+install pytorch according to https://pytorch.org/ . Anaconda is strongly recommended for example:
+```bash
+conda install pytorch
+```
+3. install
+```bash
+cd /pathto/ncnntools/pnnx
+python setup.py install
+```
+
+> **Note:**
+> If torchvision and pnnx2onnx are needed, you can set the following environment variables before 'python setup.py install' to enable them. e.g. on ubuntu:
+>
+> ```
+> export TORCHVISION_INSTALL_DIR="/project/torchvision"
+> export PROTOBUF_INCLUDE_DIR="/project/protobuf/include"
+> export PROTOBUF_LIBRARIES="/project/protobuf/lib64/libprotobuf.a"
+> export PROTOBUF_PROTOC_EXECUTABLE="/project/protobuf/bin/protoc" 
+> ```
+>
+> To do these, you must install Torchvision and Protobuf first.
+
+
+## Tests
+```bash
+cd /pathto/ncnn/tools/pnnx
+pytest python/tests/
+```
+
+## Usage
+1. export model to pnnx
+```python
+import torch
+import torchvision.models as models
+import pnnx
+
+net = models.resnet18(pretrained=True)
+x = torch.rand(1, 3, 224, 224)
+
+# You could try disabling checking when torch tracing raises error
+# mod =  pnnx.export(net, "resnet18", x, check_trace=False)
+mod = pnnx.export(net, "resnet18", x)
+```
+
+2. convert existing model to pnnx
+```python
+import pnnx
+
+pnnx.convert("resnet18.pt", [1,3,224,224], "f32")
+```
+
+## API Reference
+1. pnnx.export
+
+`model` (torch.nn.Model): model to be exported.
+
+`filename` (str): the file name.
+
+`inputs` (torch.Tensor of list of torch.Tensor) expected inputs of the model.
+
+`input_shapes` (Optional, list of int or list of list with int type inside)  shapes of model inputs.
+It is used to resolve tensor shapes in model graph. for example, [1,3,224,224] for the model with only 
+1 input, [[1,3,224,224],[1,3,224,224]] for the model that have 2 inputs. 
+
+`input_shapes2` (Optional, list of int or list of list with int type inside) shapes of alternative model inputs, 
+the format is identical to `input_shapes`. Usually, it is used with input_shapes to resolve dynamic shape (-1) 
+in model graph.
+
+`input_types` (Optional, str or list of str) types of model inputs, it should have the same length with `input_shapes`.
+for example, "f32" for the model with only 1 input, ["f32", "f32"] for the model that have 2 inputs.
+
+| typename | torch type                      |
+|:--------:|:--------------------------------|
+|   f32    | torch.float32 or torch.float    |
+|   f64    | torch.float64 or torch.double   |
+|   f16    | torch.float16 or torch.half     |
+|    u8    | torch.uint8                     |
+|    i8    | torch.int8                      |
+|   i16    | torch.int16 or torch.short      |
+|   i32    | torch.int32 or torch.int        |
+|   i64    | torch.int64 or torch.long       |
+|   c32    | torch.complex32                 |
+|   c64    | torch.complex64                 |
+|  c128    | torch.complex128                |
+
+`input_types2` (Optional, str or list of str) types of alternative model inputs.
+
+`device` (Optional, str, default="cpu") device type for the input in TorchScript model, cpu or gpu.
+
+`customop_modules` (Optional, str or list of str) list of Torch extensions (dynamic library) for custom operators.
+For example, "/home/nihui/.cache/torch_extensions/fused/fused.so" or 
+["/home/nihui/.cache/torch_extensions/fused/fused.so",...].
+
+`module_operators` (Optional, str or list of str)  list of modules to keep as one big operator. 
+for example, "models.common.Focus" or ["models.common.Focus","models.yolo.Detect"].
+
+`optlevel` (Optional, int, default=2) graph optimization level
+
+| option | optimization level                |
+|:--------:|:----------------------------------|
+|   0    | do not apply optimization         |
+|   1    | do not apply optimization         |
+|   2    | optimization more for inference   |
+
+`pnnxparam` (Optional, str, default="*.pnnx.param", * is the model name): PNNX graph definition file.
+
+`pnnxbin` (Optional, str, default="*.pnnx.bin"): PNNX model weight.
+
+`pnnxpy` (Optional, str, default="*_pnnx.py"): PyTorch script for inference, including model construction 
+and weight initialization code.
+
+`pnnxonnx` (Optional, str, default="*.pnnx.onnx"): PNNX model in onnx format.
+
+`ncnnparam` (Optional, str, default="*.ncnn.param"): ncnn graph definition.
+
+`ncnnbin` (Optional, str, default="*.ncnn.bin"): ncnn model weight.
+
+`ncnnpy` (Optional, str, default="*_ncnn.py"): pyncnn script for inference.
+
+2. pnnx.convert
+
+`ptpath` (str): torchscript model to be converted.
+
+`input_shapes` (list of int or list of list with int type inside)  shapes of model inputs.
+It is used to resolve tensor shapes in model graph. for example, [1,3,224,224] for the model with only
+1 input, [[1,3,224,224],[1,3,224,224]] for the model that have 2 inputs.
+
+`input_types` (str or list of str) types of model inputs, it should have the same length with `input_shapes`.
+for example, "f32" for the model with only 1 input, ["f32", "f32"] for the model that have 2 inputs.
+
+`input_shapes2` (Optional, list of int or list of list with int type inside) shapes of alternative model inputs,
+the format is identical to `input_shapes`. Usually, it is used with input_shapes to resolve dynamic shape (-1)
+in model graph.
+
+`input_types2` (Optional, str or list of str) types of alternative model inputs.
+
+`device` (Optional, str, default="cpu") device type for the input in TorchScript model, cpu or gpu.
+
+`customop_modules` (Optional, str or list of str) list of Torch extensions (dynamic library) for custom operators.
+For example, "/home/nihui/.cache/torch_extensions/fused/fused.so" or
+["/home/nihui/.cache/torch_extensions/fused/fused.so",...].
+
+`module_operators` (Optional, str or list of str)  list of modules to keep as one big operator.
+for example, "models.common.Focus" or ["models.common.Focus","models.yolo.Detect"].
+
+`optlevel` (Optional, int, default=2) graph optimization level
+
+`pnnxparam` (Optional, str, default="*.pnnx.param", * is the model name): PNNX graph definition file.
+
+`pnnxbin` (Optional, str, default="*.pnnx.bin"): PNNX model weight.
+
+`pnnxpy` (Optional, str, default="*_pnnx.py"): PyTorch script for inference, including model construction
+and weight initialization code.
+
+`pnnxonnx` (Optional, str, default="*.pnnx.onnx"): PNNX model in onnx format.
+
+`ncnnparam` (Optional, str, default="*.ncnn.param"): ncnn graph definition.
+
+`ncnnbin` (Optional, str, default="*.ncnn.bin"): ncnn model weight.
+
+`ncnnpy` (Optional, str, default="*_ncnn.py"): pyncnn script for inference.

tools/pnnx/python/examples/convert.py

@@ -0,0 +1,41 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import pnnx
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x):
+        x = F.relu(x)
+        return x
+
+if __name__ == "__main__":
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+
+    a0 = net(x)
+
+    mod = torch.jit.trace(net, x)
+    mod.save("test_F_relu.pt")
+
+    pnnx.convert("test_F_relu.pt", [1, 16], "f32")

tools/pnnx/python/examples/export.py

@@ -0,0 +1,41 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import pnnx
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x, y, z, w):
+        x = F.relu(x)
+        y = F.relu(y)
+        z = F.relu(z)
+        w = F.relu(w)
+        return x, y, z, w
+
+if __name__ == "__main__":
+    net = Model()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+    y = torch.rand(12, 2, 16)
+    z = torch.rand(1, 3, 12, 16)
+    w = torch.rand(1, 5, 7, 9, 11)
+
+    pnnx.export(net, "test_F_relu", (x, y, z, w))

tools/pnnx/python/pnnx/__init__.py

@@ -0,0 +1,33 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import os
+import platform
+EXEC_DIR_PATH = os.path.dirname(os.path.abspath(__file__))
+if platform.system() == 'Linux' or platform.system() == "Darwin":
+    EXEC_PATH = EXEC_DIR_PATH + "/pnnx"
+elif platform.system() == "Windows":
+    EXEC_PATH = EXEC_DIR_PATH + "/pnnx.exe"
+else:
+    raise Exception("Unsupported platform for pnnx.")
+
+from .utils.export import export
+from .utils.convert import convert
+
+try:
+    import importlib.metadata
+    __version__ = importlib.metadata.version("pnnx")
+except:
+    pass
+

tools/pnnx/python/pnnx/utils/__init__.py

@@ -0,0 +1,16 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+from .export import export
+from .convert import convert

tools/pnnx/python/pnnx/utils/convert.py

@@ -0,0 +1,94 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import os
+from .utils import check_type, generate_inputs_arg, str_in_list_to_str
+import subprocess
+from .. import EXEC_PATH
+
+def convert(ptpath, input_shapes, input_types, input_shapes2 = None,
+            input_types2 = None, device = None, customop = None,
+            moduleop = None, optlevel = None, pnnxparam = None,
+            pnnxbin = None, pnnxpy = None, pnnxonnx = None, ncnnparam = None,
+            ncnnbin = None, ncnnpy = None):
+
+    check_type(ptpath, "modelname", [str], "str")
+    check_type(input_shapes, "input_shapes", [list], "list of list with int type inside")
+    check_type(input_types, "input_types", [str, list], "str or list of str")
+    check_type(input_shapes2, "input_shapes2", [list], "list of list with int type inside")
+    check_type(input_types2, "input_types2", [str, list], "str or  list of str")
+    check_type(device, "device", [str], "str")
+    check_type(customop, "customop", [str, list], "str or list of str")
+    check_type(moduleop, "moduleop", [str, list], "str or list of str")
+    check_type(optlevel, "optlevel", [int], "int")
+
+    if input_shapes2 is None:
+        input_shapes2 = []
+    elif type(input_shapes2[0])!= list:
+        input_shapes2 = [input_shapes2]
+    if input_types2 is None:
+        input_types2 = []
+    elif type(input_types2) != list:
+        input_types2 = [input_types2]
+    if customop is None:
+        customop = []
+    elif type(customop) != list:
+        customop = [customop]
+    if moduleop is None:
+        moduleop = []
+    elif type(moduleop) != list:
+        moduleop = [moduleop]
+    if device is None:
+        device = "cpu"
+    if optlevel is None:
+        optlevel = 2
+
+    if type(input_shapes[0]) != list:
+        input_shapes = [input_shapes]
+    if type(input_types) != list:
+        input_types = [input_types]
+
+    if len(input_shapes) != len(input_types):
+        raise Exception("input_shapes should has the same length with input_types!")
+    if len(input_shapes2) != len(input_types2):
+        raise Exception("input_shapes2 should has the same length with input_types2!")
+
+    input_arg1 = generate_inputs_arg(input_shapes, input_types)
+
+    command_list = [EXEC_PATH, ptpath, "inputshape="+input_arg1,
+                    "device="+device,
+                    "optlevel="+str(optlevel)]
+    if not (len(input_shapes2) == 0):
+        input_arg2 = generate_inputs_arg(input_shapes2, input_types2)
+        command_list.append("inputshape2=" + input_arg2)
+    if not (len(customop) == 0):
+        command_list.append("customop=" + str_in_list_to_str(customop))
+    if not (len(moduleop) == 0):
+        command_list.append("moduleop=" + str_in_list_to_str(moduleop))
+    if not(pnnxparam is None):
+        command_list.append("pnnxparam="+pnnxparam)
+    if not(pnnxbin is None):
+        command_list.append("pnnxbin="+pnnxbin)
+    if not(pnnxpy is None):
+        command_list.append("pnnxpy="+pnnxpy)
+    if not(pnnxonnx is None):
+        command_list.append("pnnxonnx="+pnnxonnx)
+    if not(ncnnparam is None):
+        command_list.append("ncnnparam="+ncnnparam)
+    if not(ncnnbin is None):
+        command_list.append("ncnnbin="+ncnnbin)
+    if not(ncnnpy is None):
+        command_list.append("ncnnpy="+ncnnpy)
+    current_dir = os.getcwd()
+    subprocess.run(command_list, stdout=subprocess.PIPE, text=True, cwd=current_dir)

tools/pnnx/python/pnnx/utils/export.py

@@ -0,0 +1,169 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import os
+from .utils import check_type, get_shape_from_inputs, \
+                    get_type_from_inputs, generate_inputs_arg, str_in_list_to_str
+import subprocess
+from .. import EXEC_PATH
+
+def export(model, filename, inputs = None, input_shapes = None, input_shapes2 = None,
+           input_types = None, input_types2 = None, device = None, customop = None,
+           moduleop = None, optlevel = None, pnnxparam = None, pnnxbin = None,
+           pnnxpy = None, pnnxonnx = None, ncnnparam = None, ncnnbin = None, ncnnpy = None,
+           check_trace=True):
+    if (inputs is None) and (input_shapes is None):
+        raise Exception("inputs or input_shapes should be specified.")
+    if not (input_shapes is None) and (input_types is None):
+        raise Exception("when input_shapes is specified, then input_types should be specified correspondingly.")
+
+    check_type(filename, "filename", [str], "str")
+    check_type(inputs, "inputs", [torch.Tensor, tuple, list], "torch.Tensor or tuple/list of torch.Tensor")
+    check_type(input_shapes, "input_shapes", [list], "list of list with int type inside")
+    check_type(input_types, "input_types", [str, list], "str or list of str")
+    check_type(input_shapes2, "input_shapes2", [list], "list of list with int type inside")
+    check_type(input_types2, "input_types2", [str, list], "str or  list of str")
+    check_type(device, "device", [str], "str")
+    check_type(customop, "customop", [str, list], "str or list of str")
+    check_type(moduleop, "moduleop", [str, list], "str or list of str")
+    check_type(optlevel, "optlevel", [int], "int")
+
+    if input_shapes2 is None:
+        input_shapes2 = []
+    elif type(input_shapes2[0])!= list:
+        input_shapes2 = [input_shapes2]
+    if input_types2 is None:
+        input_types2 = []
+    elif type(input_types2) != list:
+        input_types2 = [input_types2]
+    if customop is None:
+        customop = []
+    elif type(customop) != list:
+        customop = [customop]
+    if moduleop is None:
+        moduleop = []
+    elif type(moduleop) != list:
+        moduleop = [moduleop]
+    if optlevel is None:
+        optlevel = 2
+    if type(inputs) == torch.Tensor:
+        inputs = [inputs]
+
+    if not (inputs is None):
+        model.eval()
+        mod = torch.jit.trace(model, inputs, check_trace=check_trace)
+        mod.save(filename)
+        current_path = os.path.abspath(filename)
+
+        if device is None:
+            try:
+                devicename = str(next(model.parameters()).device)
+                if ("cpu" in devicename):
+                    device = "cpu"
+                elif ("cuda" in devicename):
+                    device = "gpu"
+            except: # model without parameters
+                device = "cpu"
+
+        if input_shapes is None:
+            input_shapes = get_shape_from_inputs(inputs)
+            input_types = get_type_from_inputs(inputs)
+        else:
+            if type(input_shapes[0]) != list:
+                input_shapes = [input_shapes]
+            if type(input_types) != list:
+                input_types = [input_types]
+
+        if len(input_shapes) != len(input_types):
+            raise Exception("input_shapes should has the same length with input_types!")
+        if len(input_shapes2) != len(input_types2):
+            raise Exception("input_shapes2 should has the same length with input_types2!")
+
+        input_arg1 = generate_inputs_arg(input_shapes, input_types)
+
+        command_list = [EXEC_PATH, current_path, "inputshape=" + input_arg1,
+                        "device=" + device,
+                        "optlevel=" + str(optlevel)]
+        if not (len(input_shapes2) == 0):
+            input_arg2 = generate_inputs_arg(input_shapes2, input_types2)
+            command_list.append("inputshape2=" + input_arg2)
+        if not (len(customop) == 0):
+            command_list.append("customop=" + str_in_list_to_str(customop))
+        if not (len(moduleop) == 0):
+            command_list.append("moduleop=" + str_in_list_to_str(moduleop))
+
+        if not (pnnxparam is None):
+            command_list.append("pnnxparam=" + pnnxparam)
+        if not (pnnxbin is None):
+            command_list.append("pnnxbin=" + pnnxbin)
+        if not (pnnxpy is None):
+            command_list.append("pnnxpy=" + pnnxpy)
+        if not (pnnxonnx is None):
+            command_list.append("pnnxonnx=" + pnnxonnx)
+        if not (ncnnparam is None):
+            command_list.append("ncnnparam=" + ncnnparam)
+        if not (ncnnbin is None):
+            command_list.append("ncnnbin=" + ncnnbin)
+        if not (ncnnpy is None):
+            command_list.append("ncnnpy=" + ncnnpy)
+        current_dir = os.getcwd()
+        subprocess.run(command_list, stdout=subprocess.PIPE, text=True, cwd=current_dir)
+
+    else: # use input_shapes and input_types
+        if (input_shapes is None) or (input_types is None):
+            raise Exception("input_shapes and input_types should be specified together.")
+        model.eval()
+        mod = torch.jit.trace(model, inputs, check_trace=check_trace)
+        mod.save(filename)
+        current_path = os.path.abspath(filename)
+
+        if device is None:
+            try:
+                devicename = str(next(model.parameters()).device)
+                if ("cpu" in devicename):
+                    device = "cpu"
+                elif ("cuda" in devicename):
+                    device = "gpu"
+            except: # model without parameters
+                device = "cpu"
+
+        input_arg1 = generate_inputs_arg(input_shapes, input_types)
+
+        command_list = [EXEC_PATH, current_path, "inputshape=" + input_arg1,
+                        "device=" + device,
+                        "optlevel=" + str(optlevel)]
+        if not (len(input_shapes2) == 0):
+            input_arg2 = generate_inputs_arg(input_shapes2, input_types2)
+            command_list.append("inputshape2=" + input_arg2)
+        if not (len(customop) == 0):
+            command_list.append("customop=" + str_in_list_to_str(customop))
+        if not (len(moduleop) == 0):
+            command_list.append("moduleop=" + str_in_list_to_str(moduleop))
+        if not (pnnxparam is None):
+            command_list.append("pnnxparam=" + pnnxparam)
+        if not (pnnxbin is None):
+            command_list.append("pnnxbin=" + pnnxbin)
+        if not (pnnxpy is None):
+            command_list.append("pnnxpy=" + pnnxpy)
+        if not (pnnxonnx is None):
+            command_list.append("pnnxonnx=" + pnnxonnx)
+        if not (ncnnparam is None):
+            command_list.append("ncnnparam=" + ncnnparam)
+        if not (ncnnbin is None):
+            command_list.append("ncnnbin=" + ncnnbin)
+        if not (ncnnpy is None):
+            command_list.append("ncnnpy=" + ncnnpy)
+        current_dir = os.getcwd()
+        subprocess.run(command_list, stdout=subprocess.PIPE, text=True, cwd=current_dir)

tools/pnnx/python/pnnx/utils/utils.py

@@ -0,0 +1,91 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+
+def check_type(data, dataname, types, typesname):
+    if not(data is None):
+        if (type(data) in types):
+            return True
+        else:
+            raise Exception(dataname + " should be "+ typesname + ".")
+    else:
+        return True
+
+def get_shape_from_inputs(inputs):
+    shapes = []
+    for item in inputs:
+        sub_shapes = []
+        for l in item.shape:
+            sub_shapes.append(l)
+        shapes.append(sub_shapes)
+    return shapes
+
+def input_torch_type_to_str(tensor):
+    if tensor.dtype == torch.float32 or tensor.dtype == torch.float:
+        return "f32"
+    if tensor.dtype == torch.float64 or tensor.dtype == torch.double:
+        return "f64"
+    if tensor.dtype == torch.float16 or tensor.dtype == torch.half:
+        return "f16"
+    if tensor.dtype == torch.uint8:
+        return "u8"
+    if tensor.dtype == torch.int8:
+        return "i8"
+    if tensor.dtype == torch.int16 or tensor.dtype == torch.short:
+        return "i16"
+    if tensor.dtype == torch.int32 or tensor.dtype == torch.int:
+        return "i32"
+    if tensor.dtype == torch.int64 or tensor.dtype == torch.long:
+        return "i64"
+    if tensor.dtype == torch.complex32:
+        return "c32"
+    if tensor.dtype == torch.complex64:
+        return "c64"
+    if tensor.dtype == torch.complex128:
+        return "c128"
+
+    return "f32"
+
+def get_type_from_inputs(inputs):
+    types = []
+    for item in inputs:
+        types.append(input_torch_type_to_str(item))
+    return types
+
+def generate_inputs_arg(inputs, input_shapes):
+    generated_arg = ""
+    for i in range(0, len(inputs) - 1):
+        generated_arg += "["
+        for j in range(0, len(inputs[i]) - 1):
+            generated_arg += str(inputs[i][j]) + ','
+        generated_arg += str(inputs[i][-1])
+        generated_arg += "]"
+        generated_arg += input_shapes[i]
+        generated_arg += ","
+    generated_arg += "["
+    for j in range(0, len(inputs[-1]) - 1):
+        generated_arg += str(inputs[-1][j]) + ','
+    generated_arg += str(inputs[-1][-1])
+    generated_arg += "]"
+    generated_arg += input_shapes[-1]
+    return generated_arg
+
+def str_in_list_to_str(input_list):
+    generated_str = ""
+    for i in range(0, len(input_list) - 1):
+        generated_str += input_list[i] + ','
+    generated_str += input_list[-1]
+    return generated_str
+

tools/pnnx/python/requirements.txt

@@ -0,0 +1 @@
+torch

tools/pnnx/python/setup.py.i

@@ -0,0 +1,45 @@
+import sys
+from setuptools import setup, find_packages
+
+try:
+    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
+
+    class bdist_wheel(_bdist_wheel):
+        def finalize_options(self):
+            _bdist_wheel.finalize_options(self)
+            self.root_is_pure = False
+
+except ImportError:
+    bdist_wheel = None
+
+if sys.version_info < (3, 0):
+    sys.exit("Sorry, Python < 3.0 is not supported")
+
+requirements = ["torch"]
+
+setup(
+        name="pnnx",
+        version="${PACKAGE_VERSION}",
+        author="nihui",
+        author_email="nihuini@tencent.com",
+        description="pnnx is an open standard for PyTorch model interoperability.",
+        url="https://github.com/Tencent/ncnn/tree/master/tools/pnnx",
+        classifiers=[
+            "Programming Language :: C++",
+            "Programming Language :: Python :: 3",
+            "Programming Language :: Python :: 3.7",
+            "Programming Language :: Python :: 3.8",
+            "Programming Language :: Python :: 3.9",
+            "Programming Language :: Python :: 3.10",
+            "Programming Language :: Python :: 3.11",
+            "License :: OSI Approved :: BSD License",
+            "Operating System :: OS Independent",
+        ],
+        license="BSD-3",
+        python_requires=">=3.6",
+        packages=find_packages(),
+        package_dir={"": "."},
+        package_data={"pnnx": ["pnnx${PYTHON_MODULE_PREFIX}${PYTHON_MODULE_EXTENSION}"]},
+        install_requires=requirements,
+        cmdclass={"bdist_wheel": bdist_wheel},
+)

tools/pnnx/python/tests/test_convert.py

@@ -0,0 +1,69 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import pytest
+import pnnx
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x, y, z, w):
+        x = F.relu(x)
+        y = F.relu(y)
+        z = F.relu(z)
+        w = F.relu(w)
+        return x, y, z, w
+
+def test_convert():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+    y = torch.rand(12, 2, 16)
+    z = torch.rand(1, 3, 12, 16)
+    w = torch.rand(1, 5, 7, 9, 11)
+
+    a0, a1, a2, a3 = net(x, y, z, w)
+
+    # export torchscript
+    mod = torch.jit.trace(net, (x, y, z, w))
+    mod.save("test_F_relu_convert.pt")
+
+    pnnx.convert("test_F_relu_convert.pt",[[1,16],[12,2,16],[1,3,12,16],[1,5,7,9,11]] , ["f32", "f32", "f32", "f32"],)
+
+    # fix aten::
+    import re
+    f=open('test_F_relu_convert_pnnx.py','r')
+    alllines=f.readlines()
+    f.close()
+    f=open('test_F_relu_convert_pnnx.py','w+')
+    for eachline in alllines:
+        a=re.sub('aten::','F.',eachline)
+        a=re.sub(r'\\', r'\\\\',a)
+        f.writelines(a)
+    f.close()
+    import sys
+    import os
+    sys.path.append(os.path.join(os.getcwd()))
+    import test_F_relu_convert_pnnx
+    b0, b1, b2, b3 = test_F_relu_convert_pnnx.test_inference()
+
+    assert torch.equal(a0, b0) and torch.equal(a1, b1) and torch.equal(a2, b2) and torch.equal(a3, b3)

tools/pnnx/python/tests/test_dynamicinput_convert.py

@@ -0,0 +1,43 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import pytest
+import pnnx
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x):
+        x = F.relu(x)
+        return x
+
+def test_export():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+
+    a0 = net(x)
+
+    mod = torch.jit.trace(net, x)
+    mod.save("test_F_relu_dconvert.pt")
+
+    pnnx.convert("test_F_relu_dconvert.pt", [1, 16], "f32", input_shapes2 = [1, 8], input_types2 = "f32")

tools/pnnx/python/tests/test_dynamicinput_export.py

@@ -0,0 +1,40 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import pytest
+import pnnx
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x):
+        x = F.relu(x)
+        return x
+
+def test_export():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+
+    a0 = net(x)
+
+    pnnx.export(net, "test_F_relu_dexport", x, input_shapes2 = [1, 8], input_types2 = "f32")

tools/pnnx/python/tests/test_export.py

@@ -0,0 +1,67 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import pytest
+import pnnx
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x, y, z, w):
+        x = F.relu(x)
+        y = F.relu(y)
+        z = F.relu(z)
+        w = F.relu(w)
+        return x, y, z, w
+
+def test_export():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+    y = torch.rand(12, 2, 16)
+    z = torch.rand(1, 3, 12, 16)
+    w = torch.rand(1, 5, 7, 9, 11)
+
+    a0, a1, a2, a3 = net(x, y, z, w)
+
+    pnnx.export(net, "test_F_relu_export", (x, y, z, w))
+
+    # import sys
+    # import os
+    # sys.path.append(os.path.join(os.getcwd()))
+
+    # fix aten::
+    import re
+    f=open('test_F_relu_export_pnnx.py','r')
+    alllines=f.readlines()
+    f.close()
+    f=open('test_F_relu_export_pnnx.py','w+')
+    for eachline in alllines:
+        a=re.sub('aten::','F.',eachline)
+        a=re.sub(r'\\', r'\\\\',a)
+        f.writelines(a)
+    f.close()
+
+    import test_F_relu_export_pnnx
+    b0, b1, b2, b3 = test_F_relu_export_pnnx.test_inference()
+
+    assert torch.equal(a0, b0) and torch.equal(a1, b1) and torch.equal(a2, b2) and torch.equal(a3, b3)

tools/pnnx/python/tests/test_naiveinput_convert.py

@@ -0,0 +1,63 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import pytest
+import pnnx
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x):
+        x = F.relu(x)
+        return x
+
+def test_export():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+
+    a0 = net(x)
+
+    mod = torch.jit.trace(net, x)
+    mod.save("test_F_relu_nconvert.pt")
+
+    pnnx.convert("test_F_relu_nconvert.pt", [1, 16], "f32")
+
+    import sys
+    import os
+    sys.path.append(os.path.join(os.getcwd()))
+    # fix aten::
+    import re
+    f=open('test_F_relu_nconvert_pnnx.py','r')
+    alllines=f.readlines()
+    f.close()
+    f=open('test_F_relu_nconvert_pnnx.py','w+')
+    for eachline in alllines:
+        a=re.sub('aten::','F.',eachline)
+        a=re.sub(r'\\', r'\\\\',a)
+        f.writelines(a)
+    f.close()
+
+    import test_F_relu_nconvert_pnnx
+    b0 = test_F_relu_nconvert_pnnx.test_inference()
+
+    assert torch.equal(a0, b0)

tools/pnnx/python/tests/test_naiveinput_export.py

@@ -0,0 +1,60 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import pytest
+import pnnx
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x):
+        x = F.relu(x)
+        return x
+
+def test_export():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+
+    a0 = net(x)
+
+    pnnx.export(net, "test_F_relu_nexport", x)
+
+    import sys
+    import os
+    sys.path.append(os.path.join(os.getcwd()))
+    # fix aten::
+    import re
+    f=open('test_F_relu_nexport_pnnx.py','r')
+    alllines=f.readlines()
+    f.close()
+    f=open('test_F_relu_nexport_pnnx.py','w+')
+    for eachline in alllines:
+        a=re.sub('aten::','F.',eachline)
+        a=re.sub(r'\\', r'\\\\',a)
+        f.writelines(a)
+    f.close()
+
+    import test_F_relu_nexport_pnnx
+    b0 = test_F_relu_nexport_pnnx.test_inference()
+
+    assert torch.equal(a0, b0)

tools/pnnx/setup.py

@@ -0,0 +1,165 @@
+import io
+import os
+import sys
+import time
+import re
+import subprocess
+
+from setuptools import setup, find_packages, Extension
+from setuptools.command.build_ext import build_ext
+from setuptools.command.install import install
+
+def set_version():
+    pnnx_version = time.strftime("%Y%m%d", time.localtime())
+    return pnnx_version
+
+# Parse environment variables
+TORCH_INSTALL_DIR = os.environ.get("TORCH_INSTALL_DIR", "")
+TORCHVISION_INSTALL_DIR = os.environ.get("TORCHVISION_INSTALL_DIR", "")
+PROTOBUF_INCLUDE_DIR = os.environ.get("PROTOBUF_INCLUDE_DIR", "")
+PROTOBUF_LIBRARIES = os.environ.get("PROTOBUF_LIBRARIES", "")
+PROTOBUF_PROTOC_EXECUTABLE = os.environ.get("PROTOBUF_PROTOC_EXECUTABLE", "")
+CMAKE_BUILD_TYPE = os.environ.get("CMAKE_BUILD_TYPE", "")
+PNNX_BUILD_WITH_STATIC_CRT = os.environ.get("PNNX_BUILD_WITH_STATIC_CRT", "")
+PNNX_WHEEL_WITHOUT_BUILD = os.environ.get("PNNX_WHEEL_WITHOUT_BUILD", "")
+
+
+# Convert distutils Windows platform specifiers to CMake -A arguments
+PLAT_TO_CMAKE = {
+    "win32": "Win32",
+    "win-amd64": "x64",
+    "win-arm32": "ARM",
+    "win-arm64": "ARM64",
+}
+
+# A CMakeExtension needs a sourcedir instead of a file list.
+# The name must be the _single_ output extension from the CMake build.
+# If you need multiple extensions, see scikit-build.
+class CMakeExtension(Extension):
+    def __init__(self, name, sourcedir=""):
+        Extension.__init__(self, name, sources=[])
+        self.sourcedir = os.path.abspath(sourcedir)
+
+class CMakeBuild(build_ext):
+    def build_extension(self, ext):
+        extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
+        extdir = os.path.join(extdir, "pnnx")
+
+        # required for auto-detection of auxiliary "native" libs
+        if not extdir.endswith(os.path.sep):
+            extdir += os.path.sep
+
+        cfg = "Debug" if self.debug else "Release"
+
+        # CMake lets you override the generator - we need to check this.
+        # Can be set with Conda-Build, for example.
+        cmake_generator = os.environ.get("CMAKE_GENERATOR", "")
+
+        # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON
+        # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code
+        # from Python.
+        cmake_args = [
+            "-DCMAKE_RUNTIME_OUTPUT_DIRECTORY={}".format(extdir),
+            "-DCMAKE_RUNTIME_OUTPUT_DIRECTORY_RELEASE={}".format(extdir),
+            "-DPython3_EXECUTABLE={}".format(sys.executable),
+            "-DCMAKE_BUILD_TYPE={}".format(cfg),  # not used on MSVC, but no harm
+        ]
+
+        if TORCH_INSTALL_DIR != "":
+            cmake_args.append("-DTorch_INSTALL_DIR=" + TORCH_INSTALL_DIR)
+        if TORCHVISION_INSTALL_DIR != "":
+            cmake_args.append("-DTorchVision_INSTALL_DIR=" + TORCHVISION_INSTALL_DIR)
+        if PROTOBUF_INCLUDE_DIR != "":
+            cmake_args.append("-DProtobuf_INCLUDE_DIR=" + PROTOBUF_INCLUDE_DIR)
+        if PROTOBUF_LIBRARIES != "":
+            cmake_args.append("-DProtobuf_LIBRARIES=" + PROTOBUF_LIBRARIES)
+        if PROTOBUF_PROTOC_EXECUTABLE != "":
+            cmake_args.append("-DProtobuf_PROTOC_EXECUTABLE=" + PROTOBUF_PROTOC_EXECUTABLE)
+        if CMAKE_BUILD_TYPE != "":
+            cmake_args.append("-DCMAKE_BUILD_TYPE=" + CMAKE_BUILD_TYPE)
+        if PNNX_BUILD_WITH_STATIC_CRT != "":
+            cmake_args.append("-DPNNX_BUILD_WITH_STATIC_CRT=" + PNNX_BUILD_WITH_STATIC_CRT)
+            
+        build_args = []
+
+        if self.compiler.compiler_type == "msvc":
+            # Single config generators are handled "normally"
+            single_config = any(x in cmake_generator for x in {"NMake", "Ninja"})
+
+            # CMake allows an arch-in-generator style for backward compatibility
+            contains_arch = any(x in cmake_generator for x in {"ARM", "Win64"})
+
+            # Specify the arch if using MSVC generator, but only if it doesn't
+            # contain a backward-compatibility arch spec already in the
+            # generator name.
+            if not single_config and not contains_arch:
+                cmake_args += ["-A", PLAT_TO_CMAKE[self.plat_name]]
+
+            # Multi-config generators have a different way to specify configs
+            if not single_config:
+                cmake_args += [
+                    "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(cfg.upper(), extdir)
+                ]
+                build_args += ["--config", cfg]
+
+        # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level
+        # across all generators.
+        if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ:
+            # self.parallel is a Python 3 only way to set parallel jobs by hand
+            # using -j in the build_ext call, not supported by pip or PyPA-build.
+            if hasattr(self, "parallel") and self.parallel:
+                # CMake 3.12+ only.
+                build_args += ["-j{}".format(self.parallel)]
+            else:
+                build_args += ["-j2"]
+
+        if not os.path.exists(self.build_temp):
+            os.makedirs(self.build_temp)
+
+        if not (PNNX_WHEEL_WITHOUT_BUILD == "ON"):
+            subprocess.check_call(
+                ["cmake", ext.sourcedir] + cmake_args, cwd=self.build_temp
+            )
+            subprocess.check_call(
+                ["cmake", "--build", "."] + build_args, cwd=self.build_temp
+            )
+        else:
+            pass
+
+if sys.version_info < (3, 0):
+    sys.exit("Sorry, Python < 3.0 is not supported")
+
+requirements = ["torch"]
+
+with io.open("README.md", encoding="utf-8") as h:
+    long_description = h.read()
+
+setup(
+    name="pnnx",
+    version=set_version(),
+    author="nihui",
+    author_email="nihuini@tencent.com",
+    description="pnnx is an open standard for PyTorch model interoperability.",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/Tencent/ncnn/tree/master/tools/pnnx",
+    classifiers=[
+        "Programming Language :: C++",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "License :: OSI Approved :: BSD License",
+        "Operating System :: OS Independent",
+    ],
+    license="BSD-3",
+    python_requires=">=3.7",
+    packages=find_packages("python"),
+    package_data={"pnnx": ["pnnx", "pnnx.exe"]},
+    package_dir={"": "python"},
+    install_requires=requirements,
+    ext_modules=[CMakeExtension("pnnx")],
+    cmdclass={"build_ext": CMakeBuild},
+)