|
目录
前言?
1.torch下将模型转换为onnx模型
2.实际演示转换
3.使用
4.结尾
前言?
ONNX（Open Neural Network Exchange）是一种开放式的文件格式，可以用于保存不同深度学习框架下的网络模型和参数，从而方便将模型进行不同框架下的转换。
1.torch下将模型转换为onnx模型
这里介绍一个函数——torch.onnx.export()：
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None)
其中主要使用的参数含义如下：
model(torch.nn.Module)?:需要转换的模型。
args(tuple or torch.Tensor)?:模型的输入。因为expoort需要运行模型，所以需要为其提供一个输入。输入的尺寸一致，输入的值随机即可。（注意，除非使用动态轴，否则onnx模型的输入尺寸将被固定为该输入的尺寸）
f:保存路径
export_params(bool)：默认为True。如果默认或设置为True，则会将导入模型的权重参数正常保存。如果为False，则会保存一个未训练过的模型。
verbose(bool):如果为True，则打印被导出到标准输出的模型的描述。
input_names(list of str): 按顺序分配名称到中的输入节点。
output_names(list of str):按顺序分配名称到图中的输出节点。
dynamic_axes(dict<string, dict<int, string>> or dict<string, list(int)>, default empty dict)：默认情况下，导出的模型将有所有输入和输出张量的形状设置为完全匹配args中给出的形状。要指定张量的轴为动态的(即只有在运行时才知道)，要将dynamic_axes设置为一个带schema的字典，如下形式:
torch.onnx.export(SumModule(), (torch.ones(2, 2),), "onnx.pb",
input_names=["x"], output_names=["sum"],
dynamic_axes={
# dict value: manually named axes
"x": {0: "my_custom_axis_name"},
# list value: automatic names
"sum": [0],
})
2.实际演示转换
这里我写了一段深度学习的代码，以此为例来进行实际演示onnx模型的使用。代码如下：
import torch
from torch import nn
from torch.autograd import Variable
class lstm_reg(nn.Module):
def __init__(self, input_size, hidden_size, output_size=1, num_layers=2):
super(lstm_reg, self).__init__()
self.rnn = nn.LSTM(input_size, hidden_size, num_layers,device='cpu') # rnn
self.reg = nn.Linear(hidden_size, output_size,device='cpu') # 回归
def forward(self, x):
x, _ = self.rnn(x)
s, b, h = x.shape
x = x.view(s * b, h)
x = self.reg(x)
x = x.view(s, b, -1)
return x
#随机网络的输入和输出，输入是1000条1x12的向量，输出是1000个值
device=torch.device('cuda')
x=torch.randn(1000,1,12)
x=torch.as_tensor(x,dtype=torch.float64)
y=torch.randn(1000,1,1)
y=torch.as_tensor(y,dtype=torch.float64)
#对这一千条数据划分，七成是训练集，三成是测试集
train_size=int(len(x)*0.7)
test_size=len(x)-train_size
train_x,train_y=x[0:train_size,:,:],y[0:train_size,:,:]
test_x,test_y=x[train_size:,:,:],y[train_size:,:,:]
train=torch.utils.data.TensorDataset(train_x,train_y)
val=torch.utils.data.TensorDataset(test_x,test_y)
net = lstm_reg(12, 100)#net为我们定义的网络,12是输入向量的长度，100是lstm层的深度
net=net.to(device)
criterion = nn.MSELoss()#均方损失函数
optimizer = torch.optim.Adam(net.parameters(), lr=1e-2)#梯度优化方法
if __name__ == '__main__':
for e in range(1):
net=net.train()
train_loader = torch.utils.data.DataLoader(train,
batch_size=32,
num_workers=0,
shuffle=True,
pin_memory=True,
)
val_loader = torch.utils.data.DataLoader(val,
batch_size=4,
num_workers=0,
pin_memory=True,
)
for step,(v_x,v_y)in enumerate(train_loader):
v_x = Variable(v_x).float()
v_y = Variable(v_y).float()
v_y=v_y.cuda()
v_x=v_x.cuda()
# 前向传播
out = net(v_x)
loss = criterion(out, v_y)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Epoch: {},step:{} Loss: {:.5f}'.format(e + 1,step,loss.item()))
#本来这里想再写个测试，但是又觉得没有必要...
torch.save(net, "model.pt")
这个代码网络结构为lstm加一层线性层。输入数据和目标输出用了随机的两个张量。
将这段代码运行之后，我们可以得到一个“model.pt”的文件，这是我们在用pytorch进行训练时常常会见到的一种文件格式。下一步我们要做的就是将这个模型转换为onnx文件。代码如下：
import torch
model=torch.load("model.pt",map_location="cpu")#加载需要转换的模型
dummy_input = torch.randn(1,1,12)#随机一个输入
input_names = ["input"]#输入名称
output_names = ["output"]#输出名称
torch.onnx.export(model, dummy_input, r"model.onnx", verbose=False, input_names=input_names, output_names=output_names)
需要注意，在执行torch.save(net, "model.pt")时，保存的是整个网络模型——包括网络框架和权重参数。而大部分时候，为了更灵活对待训练好的模型参数，常常只保存权重参数，使用torch.save(net.state_dict(), 'model.pt')进行保存。这种情况下是不能直接使用上述代码进行模型转换的。
当只保存了网络的权重参数时，我们先要重新加载模型，并导入网络的权重参数。现在我们先将训练代码的最后一行改为，仅保存权重参数：
torch.save(net.state_dict(), "model_only-weights.pt")
随后， 首先从我们定义的lstm_reg中重新加载模型，再进行模型转换，代码如下：
import torch
import demo
from demo import lstm_reg
model=lstm_reg(12,100)#加载模型
weight=torch.load("model_only-weights.pt",map_location="cpu")#加载权重参数
model.load_state_dict(weight)#将权重参数导入模型
dummy_input = torch.randn(1,1,12)
input_names = ["input"]
output_names = ["output"]
torch.onnx.export(model, dummy_input, r"model.onnx", verbose=False, input_names=input_names, output_names=output_names,)
3.使用
这里需要用到一个python库——onnxruntime。直接使用pip install onnxruntime下载即可。使用onnx模型推理代码如下：
import torch
import numpy
import onnxruntime
dummy_input = torch.randn(1,1,12)
session = onnxruntime.InferenceSession("model.onnx")
result = session.run([], {"input":dummy_input.numpy()})
print(result)
输出结果为：
[array([[[-0.09918265]]], dtype=float32)]
4.结尾
至此，onnx模型的转换和简单的使用已经讲完了。祝愿大家共同进步，技术越来越强。?
|