pytorch导出onnx简单易懂(固定输入尺寸以及动态输入尺寸)

先简单定义一个网络结构

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.ClassifierLayer = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=4, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=4, out_channels=4, kernel_size=3, stride=1),
            nn.MaxPool2d(kernel_size=2),
        )


    def forward(self, x):
        y = self.ClassifierLayer(x)
        y = F.interpolate(y,size=x.shape[2:])
        return y

转换代码(固定输入尺寸)

#实例化网络结构
model = Net()
#定义一个模拟的输入
input_data = torch.randn([1,3,320,320])
#参数1:模型  参数2:模拟的输入数据  参数3:导出的onnx绝对路径  #参数4:是否打印导出过程中的详细信息
torch.onnx.export(model, input_data, "Net.onnx", verbose=True)

接下来可视化一下导出后的onnx网络结构
网络结构图
注意看右边的输入部分[1,3,320,320]

动态输入尺寸转换代码

    input_name = 'input'
    output_name = 'output'
    torch.onnx.export(model, 
                      input_data, 
                      "simplenet.onnx",
                      opset_version=11,
                      input_names=[input_name],
                      output_names=[output_name],
                      dynamic_axes={
                          input_name: {0: 'batch_size', 2: 'in_width', 3: 'int_height'},
                          output_name: {0: 'batch_size', 2: 'out_width', 3: 'out_height'}}
                      )

网络结构图
同样注意右边的输入部分,现在变成了[batch_size,3,in_width,int_height]

至此最基础的pytorch转onnx就全部完成了,所有的网络结构都可套用这两段代码,不过如果遇到了转换失败算子不支持的问题,可以尝试修改opset_version参数,或者使用其他结构去代替不支持的操作

再补充一下onnxsim:
如果你想把onnx再转其他的模型,就拿ncnn来举例,上图中的Gather在转换过程中就会出现不支持转到ncnn的错误信息。
推荐使用onnxsim,不论任何模型都可以用,当然能否sim成功还要看你的模型复杂程度,不过目前遇到的大部分模型都是可以经过sim规避很多问题的。
项目地址:onnxsim
假设你已经安装好了onnxsim,那么

onnxsim xxx.onnx xxx-sim.onnx

经过onnxsim之后,你的onnx模型就会简洁很多了,以静态输入尺寸举例如下图
请添加图片描述
可以看到ncnn不支持的Gather层已经被简化掉了。

注意:onnxsim的用法这里只说明了最基本的,其余的更多用法自行查看参数含义去使用。

<think>好的,用户想了解如何将PyTorch部署到硬件上。首先,我需要回忆PyTorch部署的相关知识,确保信息准确。可能涉及模型优化、不同硬件平台的适配等步骤。 首先,用户可能是开发者或研究人员,他们可能已经训练好了模型,现在需要在实际环境中运行。需要明确部署的流程,可能包括模型转换、优化、硬件选择等。 接下来,应该分步骤讲解。第一步通常是模型准备,比如导出模型为TorchScript或ONNX格式。然后是硬件选择,不同的硬件可能需要不同的工具链,比如NVIDIA的TensorRT,移动端用Core ML或TensorFlow Lite。 需要考虑用户可能的困惑点,比如如何优化模型性能,如何处理不同硬件上的兼容性问题。需要提到量化、剪枝等方法,以及常用的部署框架如LibTorch或ONNX Runtime。 另外,硬件类型多样,比如边缘设备、嵌入式系统、服务器等,可能需要分别说明。例如,部署到嵌入式设备可能需用NCNN或TFLite Micro,而服务器端可能用Docker容器化。 需要检查是否有遗漏的关键点,比如实时性要求、资源限制等,这些可能影响部署策略。还要确保提到的工具和库是当前主流的,避免过时信息。 最后,总结步骤,强调测试和性能优化的重要性,确保用户理解整个流程。避免使用过于专业的术语,保持解释清晰易懂。</think>将PyTorch模型部署到硬件设备需要经过模型优化、格式转换和硬件适配等步骤。以下是分步说明: 1. **模型准备** - 使用`torch.jit.trace`或`torch.jit.script`导出为TorchScript格式 ```python traced_model = torch.jit.trace(model, example_input) traced_model.save("model.pt") ``` - 或导出ONNX格式实现跨平台兼容: ```python torch.onnx.export(model, dummy_input, "model.onnx") ``` 2. **硬件平台选择** - **移动端**:通过Core ML(iOS)或TensorFlow Lite(Android)转换 - **嵌入式设备**:使用NCNN/TFLite Micro量化部署 - **GPU服务器**:通过TensorRT加速 - **浏览器端**:转换为WebAssembly格式 3. **性能优化技术** - **量化**:8-bit量化减小模型体积 ```python quantized_model = torch.quantization.quantize_dynamic(model, dtype=torch.qint8) ``` - **剪枝**:移除冗余参数 - **内核融合**:优化计算图结构 4. **部署工具链** - **LibTorch**:C++推理接口 ```cpp torch::jit::Module module = torch::jit::load("model.pt"); ``` - **ONNX Runtime**:跨平台推理引擎 - **TVM**:自动生成硬件优化代码 5. **硬件特定优化** - NVIDIA GPU:启用TensorCore加速 - ARM芯片:使用NEON指令集优化 - FPGA:通过OpenCL重写计算内核 6. **部署验证** - 精度测试:比对输出误差 - 性能测试:测量FPS和延迟 - 内存分析:监控显存/内存占用 典型部署流程示例: 1. PyTorch模型 -> ONNX导出 -> TensorRT优化 -> NVIDIA Jetson部署 2. PyTorch模型 -> TorchScript -> LibTorch嵌入 -> 工业控制器 3. PyTorch模型 -> TVM编译 -> RISC-V芯片部署 注意事项: - 不同硬件平台对算子支持存在差异 - 部署前需测试目标设备的指令集兼容性 - 实时系统需考虑确定性执行要求 - 边缘设备部署需优化功耗和散热 最新趋势: - 使用TorchServe进行容器化部署 - 通过ONNX Runtime实现异构计算 - 采用Apache TVM进行自动硬件代码生成
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

木辛叔叔

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值