onnx学习

模型导出测试

import numpy as np
import torch
import torch.nn as nn

## 构造随机输入，参数为shape，输入满足正太分布
input = torch.randn(4, 256, 256)
input=torch.randint_like(input, low=0, high=8)
print(input)

# 将测试数据保存成npy文件
testdata = input.detach().numpy()
np.save('./input.npy', testdata)

# 将测试数据保存成txt文件
testdata = testdata.reshape(-1)
np.savetxt('./input.txt',testdata, fmt="%s",delimiter='\n',newline='\n')


# 增加batch参数（相当于只有一张图片）
input = input.reshape(1,4,256,256)
# With square kernels and equal stride
m = nn.Conv2d(4, 24, 3, stride=2,padding=1)
m.weight.data = torch.randint_like(m.weight.data, low=0, high=8)
m.bias.data=torch.ones_like(m.bias.data)
output = m(input)
print(output.shape)


# 保存模型预测后的数据
arr = output.detach().numpy()
np.save('./output.npy', arr)
arr = arr.reshape(-1)
np.savetxt('./output.txt',arr, fmt="%s",delimiter='\n',newline='\n')

# 保存模型
torch.save(m.state_dict(), "./torch_model.pth")
m.load_state_dict(torch.load("./torch_model.pth", map_location='cpu'))
dummy_input = torch.randint_like(input, low=0, high=8)
torch.onnx.export(m, dummy_input, "./onnx_model.onnx", verbose=True)

模型输出测试

import torch
import onnxruntime
import numpy as np


# 构建输入数据方式一:加载npy文件
im = np.load("test_input.npy")
# 构建输入数据方式二：加载txt文件
'''
arr = np.loadtxt('test_input.txt')
im = np.reshape(arr,(3,640,640))
'''
# 构建输入数据方式三：numpy构建数组
'''
# 构建全1数组
im = np.ones([3,640,640], dtype = int)
np.save("filename.npy",im)
'''

im = torch.from_numpy(im).to('cpu')
im = im.float()
# 数据归一化，根据需要进行修改
im /= 255  # 0 - 255 to 0.0 - 1.0
# 增加batch纬度
if len(im.shape) == 3:
    im = im[None]  # expand for batch dim

onnx_path = "./onnx.onnx"
ort_session = onnxruntime.InferenceSession(onnx_path)


data = im.numpy()
ort_inputs = {ort_session.get_inputs()[0].name: data}
ort_outs = ort_session.run(None, ort_inputs)
torch.set_printoptions(precision=4, sci_mode=False)
np.set_printoptions(precision=4, suppress=True)
print("onnx_result: ", len(ort_outs))
print("onnx_result: ", ort_outs)