模型部署 - onnx 的导出和分析 - PyTorch 导出 ONNX - 学习记录

onnx 的导出和分析

  • 一、PyTorch 导出 ONNX 的方法
    • 1.1、一个简单的例子 -- 将线性模型转成 onnx
    • 1.2、导出多个输出头的模型
    • 1.3、导出含有动态维度的模型
  • 二、pytorch 导出 onnx 不成功的时候如何解决
    • 2.1、修改 opset 的版本
    • 2.2、替换 pytorch 中的算子组合
    • 2.3、在 pytorch 登记( 注册 ) onnx 中某些算子
      • 案例一:
      • 2.3.1、注册方法一
      • 2.3.2、注册方法二
      • 案例二:
    • 2.4、直接修改 onnx,创建 plugin

一、PyTorch 导出 ONNX 的方法

1.1、一个简单的例子 – 将线性模型转成 onnx

首先我们用 pytorch 定义一个线性模型,nn.Linear : 线性层执行的操作是 y = x * W^T + b,其中 x 是输入,W 是权重,b 是偏置。(实际上就是一个矩阵乘法)

class Model(torch.nn.Module):
    def __init__(self

你可能感兴趣的:(Onnx与TensorRT,pytorch,学习,人工智能)