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