-
Notifications
You must be signed in to change notification settings - Fork 4.4k
use ncnn with pytorch or onnx
This guide is designed to help pytorch and onnx users use the new-generation model conversion tool, pnnx, to efficiently and reliably convert models to the ncnn format for high-performance inference on the edge.
This document is written and revised based on the official pnnx documentation.
- pnnx project: https://github.com/pnnx/pnnx
- ncnn project: https://github.com/Tencent/ncnn
- supported pytorch operators: https://github.com/Tencent/ncnn/tree/master/tools/pnnx#supported-pytorch-operator-status
- supported onnx operators: https://github.com/Tencent/ncnn/tree/master/tools/pnnx#supported-onnx-operator-status
Regardless of which framework you come from, pnnx offers significant advantages over traditional tools (like onnx2ncnn):
-
Forget the Hassles of onnx: The traditional
pytorch -> onnx -> ncnnpipeline often fails due to onnx operator compatibility issues and dynamic shape problems. pnnx can convert directly from pytorch, completely bypassing the unstable intermediate step of onnx. - Core Framework Support: pnnx focuses on supporting pytorch and onnx, providing you with a unified and consistent conversion experience.
- More Stable and Powerful: pnnx can handle a wider range of modern operators and complex model architectures, generating cleaner and more accurate ncnn graphs.
- Active and Continuous Development: pnnx is under active development, constantly adding support for the latest operators and features from both source frameworks and the ncnn engine.
- Richer Graph Information: pnnx preserves the original model's structural information during the conversion process, which is highly beneficial for model analysis and subsequent optimization.
For pytorch users, converting directly from a pytorch model is the most stable and efficient path.
This is the simplest and most recommended workflow, allowing you to complete the model conversion with a single command without leaving your Python environment.
First, install the pnnx Python package. This command installs both the pnnx Python library and the pnnx command-line tool.
pip3 install pnnxCalling the pnnx.export function will generate both a TorchScript (.pt) file and the .param and .bin files required by ncnn.
Complete Code Example:
import torch
import torch.nn as nn
import pnnx
# 1. Define or load your pytorch model
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3, 1, 1)
self.relu = nn.ReLU()
self.fc = nn.Linear(16 * 224 * 224, 10)
def forward(self, x):
x = self.conv1(x)
x = self.relu(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
# 2. Instantiate the model and set it to evaluation mode
model = MyModel()
model.eval()
# 3. Create a dummy input tensor with the correct input shape
input_tensor = torch.rand(1, 3, 224, 224)
# 4. Call pnnx.export to export the model
pnnx.export(model, "my_model.pt", (input_tensor,))
print("Conversion complete!")
print("Please check for the generated my_model.pt, my_model.ncnn.param, and my_model.ncnn.bin files.")If you have already run pip install pnnx, the pnnx command is available, and you can proceed to the next step.
For non-Python environments or users who prefer a standalone executable, you can manually download the latest binary from the pnnx Releases page.
import torch
# ... (model definition from above)
model = MyModel()
model.eval()
input_tensor = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, input_tensor)
traced_script_module.save("my_model.pt")Run the following command in your terminal.
# Syntax: pnnx <torchscript_model_path>
pnnx my_model.ptFor users who already have an .onnx file, please use pnnx for conversion.
-
Method 1 (Recommended): If you have Python in your environment, install it directly via pip.
pip3 install pnnx
The
pnnxcommand will be automatically added to your system's path. -
Method 2 (Alternative): For non-Python environments or to use a standalone program, you can download the latest executable from the pnnx Releases page.
-
Method 3 (Convert Everywhere with GUI): To convert models to ncnn anytime and anywhere, you can use pnnx.js. To protect your data security, the website is built based on wasm64 technology, and your models will not be uploaded to the server.
Open a terminal, navigate to the directory containing your model file, and run the following command.
Basic Command Example:
# Syntax: pnnx <onnx_model_path>
pnnx my_model.onnxAfter the command executes successfully, you will get the my_model.ncnn.param and my_model.ncnn.bin files, which can be directly loaded and used in your ncnn project.