diff --git a/.gitignore b/.gitignore index aa8ea4ddcb2..2c71aee0332 100644 --- a/.gitignore +++ b/.gitignore @@ -54,6 +54,7 @@ __pycache__ *.pyd *.egg-info/ python/setup.py +tools/pnnx/python/setup.py # Clangd .cache/ diff --git a/tools/pnnx/python/README.md b/tools/pnnx/python/README.md new file mode 100644 index 00000000000..48241f13381 --- /dev/null +++ b/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. diff --git a/tools/pnnx/python/examples/convert.py b/tools/pnnx/python/examples/convert.py new file mode 100644 index 00000000000..ab695186bd6 --- /dev/null +++ b/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") diff --git a/tools/pnnx/python/examples/export.py b/tools/pnnx/python/examples/export.py new file mode 100644 index 00000000000..c31d9046c59 --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/tools/pnnx/python/pnnx/__init__.py b/tools/pnnx/python/pnnx/__init__.py new file mode 100644 index 00000000000..0ba69277d7a --- /dev/null +++ b/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 + diff --git a/tools/pnnx/python/pnnx/utils/__init__.py b/tools/pnnx/python/pnnx/utils/__init__.py new file mode 100644 index 00000000000..e83f4bc373a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tools/pnnx/python/pnnx/utils/convert.py b/tools/pnnx/python/pnnx/utils/convert.py new file mode 100644 index 00000000000..48c8f3f8fd6 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/tools/pnnx/python/pnnx/utils/export.py b/tools/pnnx/python/pnnx/utils/export.py new file mode 100644 index 00000000000..21cb5c3af54 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/tools/pnnx/python/pnnx/utils/utils.py b/tools/pnnx/python/pnnx/utils/utils.py new file mode 100644 index 00000000000..1af06114be8 --- /dev/null +++ b/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 + diff --git a/tools/pnnx/python/requirements.txt b/tools/pnnx/python/requirements.txt new file mode 100644 index 00000000000..08ed5eeb4b9 --- /dev/null +++ b/tools/pnnx/python/requirements.txt @@ -0,0 +1 @@ +torch \ No newline at end of file diff --git a/tools/pnnx/python/setup.py.i b/tools/pnnx/python/setup.py.i new file mode 100644 index 00000000000..5c3c72e876d --- /dev/null +++ b/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}, +) \ No newline at end of file diff --git a/tools/pnnx/python/tests/test_convert.py b/tools/pnnx/python/tests/test_convert.py new file mode 100644 index 00000000000..23e1d0d4a90 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/tools/pnnx/python/tests/test_dynamicinput_convert.py b/tools/pnnx/python/tests/test_dynamicinput_convert.py new file mode 100644 index 00000000000..8ddb9a3e372 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/tools/pnnx/python/tests/test_dynamicinput_export.py b/tools/pnnx/python/tests/test_dynamicinput_export.py new file mode 100644 index 00000000000..3436b07eba5 --- /dev/null +++ b/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") diff --git a/tools/pnnx/python/tests/test_export.py b/tools/pnnx/python/tests/test_export.py new file mode 100644 index 00000000000..26d8bc41fd5 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/tools/pnnx/python/tests/test_naiveinput_convert.py b/tools/pnnx/python/tests/test_naiveinput_convert.py new file mode 100644 index 00000000000..af4a53108c4 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/tools/pnnx/python/tests/test_naiveinput_export.py b/tools/pnnx/python/tests/test_naiveinput_export.py new file mode 100644 index 00000000000..d795ba3c743 --- /dev/null +++ b/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) diff --git a/tools/pnnx/setup.py b/tools/pnnx/setup.py new file mode 100644 index 00000000000..dcbb99922b3 --- /dev/null +++ b/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}, +)