ONNX实践

有关ONNX的大体介绍已经很多了，比如ONNX简明教程，但相关的代码要么很复杂，要么太简单，这里将介绍三个方面：

• 如何生成.onnx
• 如何查看.onnx
• 如何转成.json格式

一、如何生成.onnx

import torch
import torch.nn as nn
import torch.nn.functional as F
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1,out_channels=6,kernel_size=3,stride=1,padding=0)
        self.conv2 = nn.Conv2d(in_channels=6,out_channels=16,kernel_size=3,stride=1,padding=0)
        self.fc1   = nn.Linear(16*5*5, 120)
        self.fc2   = nn.Linear(120, 84)
        self.fc3   = nn.Linear(84, 10)

    def forward(self, x):
        out = self.conv1(x)# 26x26
        out = F.max_pool2d(F.relu(out), 2) # 13x13
        out = self.conv2(out) #11x11
        out = F.max_pool2d(F.relu(out), 2)  # 5x5
        out = out.view(out.size(0), -1) #400
        out = self.fc1(out) # 120
        out = self.fc2(F.relu(out)) # 84
        out = self.fc3(F.relu(out))  #10
        return out
# 1、加载模型之前训练好的模型
net = LeNet()
net = net.to('cpu')
checkpoint = torch.load('LeNet_mnist.pth')
net.load_state_dict(checkpoint['net'])
# 2、导出
batch_size = 1  
input_shape = (1,28,28)   
x = torch.randn(batch_size,*input_shape)# 一张图片输入
net.eval()
torch.onnx.export(net,x,"./mnist.onnx")

LeNet_mnist.zip解压后要放在当前目录，运行代码就可以生成mnist.onnx文件了。

二、如何查看计算图？

可以通过一个网站Netron，open Model，选择刚刚生成的mnist.onnx文件即可。不但可以查看网络结构，而且包含权重信息。

三、如何转.json?

计算图后期有其他用途，希望导成可读的文件，.json是很好的格式，找了很久才搜到转换的方式：

import onnx
import json
from google.protobuf.json_format import MessageToJson

onnx_model = onnx.load("./mnist.onnx")
s = MessageToJson(onnx_model)
onnx_json = json.loads(s)

output_json_path = 'mnist2.json'

with open(output_json_path, 'w') as f:
    json.dump(onnx_json, f, indent=2)

就这简单的几句，花了我不少时间onnx2json，最后通过文本编辑器或者Firefox浏览器可以产看其结构，但其权重似乎以二进制文件保存。