【推理引擎】ONNX 模型解析

Python微信订餐小程序课程视频

https://edu.csdn.net/course/detail/36074

Python实战量化交易理财系统

https://edu.csdn.net/course/detail/35475

定义模型结构

首先使用 PyTorch 定义一个简单的网络模型：

class ConvBnReluBlock(nn.Module):
    def \_\_init\_\_(self) -> None:
        super().__init__()

        self.conv1 = nn.Conv2d(3, 64, 3)
        self.bn1 = nn.BatchNorm2d(64)
        self.maxpool1 = nn.MaxPool2d(3, 1)

        self.conv2 = nn.Conv2d(64, 32, 3)
        self.bn2 = nn.BatchNorm2d(32)
    
        self.relu = nn.ReLU()

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.maxpool1(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        
        return out

在导出模型之前，需要提前定义一些变量：

model = ConvBnReluBlock()     # 定义模型对象
x = torch.randn(2, 3, 255, 255)      # 定义输入张量

然后使用 PyTorch 官方 API（torch.onnx.export）导出 ONNX 格式的模型：

# way1:
torch.onnx.export(model, (x), "conv\_bn\_relu\_evalmode.onnx", input_names=["input"], output_names=['output'])

# way