pytorch训练出来的模型参数如何保存为嵌入式C语言能够调用的形式

在深度学习算法嵌入式开发应用中，如何把在pytorch下训练得到的模型参数，提取出来保存成C语言能够调用的数组形式，是嵌入式应用开发的重要步骤。

以下介绍一个简单的例子，下文中的代码来自莫凡教程，是一个线性回归的简单深度学习例子，输入节点为1个，隐藏层10个节点，输出一个节点。relu用的是普通的，所以不需要参数。

 

从模型中索引到参数通过如下方法调用named_parameters()，parameters.detach().numpy()

for name,parameters in net3.named_parameters():
    print(name,':',parameters.size())
    parm[name] = parameters.detach().numpy()

 

保存成文件如下操作：

f = open('./w1.txt',mode = 'w+') #是创建一个txt文件。参数 'w+'的意思是如果存在该文件，则删除内容，从头开始书写，如果不存在该文件，则创建一个。

f.write('const float w1[%d] = {\n'% w1.shape[0]*w1.shape[1]) #这一句的意思是定义一个float 类型的数组名字为w1的常量，数组的维度 = w1.shape[0]*w1.shape[1]。

for j in range(w1.shape[0]):#遍历每一行（=10）
    for i in range(w1.shape[1]):#遍历每一列（=1）
        f.write(str(w1[j,i])+',\n')#由于这种形式可以操控任何一个参数，所以你可以想保存成怎样就怎样。str（）可以直接把numpy类型的数转成str。因为f.write（）支持的是str类型。当然还需要加点                                             #逗号、回车之类以写成c语言认识的形式
        pass
    pass

 

最终输出是这样的文件

w1.txt

const float w1[10] = {
-1.016925,
-0.3459752,
0.5006739,
-0.60270876,
0.4570868,
0.025785923,
0.154984,
0.3099202,
0.6377684,
-0.55719304,
}

 

代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import time

x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)  # x data (tensor), shape=(100, 1)
y = x.pow(2) + 0.2*torch.rand(x.size())                 # noisy y data (tensor), shape=(100, 1)

class Net(nn.Module):
    def __init__(self,n_feature,n_hidden,n_output):
        super(Net,self).__init__()
        self.hidden = nn.Linear(n_feature,n_hidden)
        self.predict = nn.Linear(n_hidden,n_output)
        
    def forward(self,x):
        x = F.relu(self.hidden(x))
        x = self.predict(x)
        
        return x

net1 = Net(1,10,1)
optimizer = torch.optim.SGD(net1.parameters(),lr=0.1)
loss_func = nn.MSELoss()

for i in range(100):
    prediction = net1(x)
    loss = loss_func(prediction,y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    

torch.save(net1,'net.pkl')
torch.save(net1.state_dict(),'net_parms.pkl')

net2 = torch.load('net.pkl')
net3 = Net(1,10,1)
net3.load_state_dict(torch.load('net_parms.pkl'))

parm = {}
for name,parameters in net3.named_parameters():
    print(name,':',parameters.size())
    parm[name] = parameters.detach().numpy()

w1 = parm['hidden.weight']
b1 = parm['hidden.bias']
w2 = parm['predict.weight']
b2 = parm['predict.bias']
print(w1)
print(b1)
print(w2)
print(b2)
# print('size:',w1.shape[0])
f = open('./w1.txt',mode = 'w+')
f.write('const float w1[%d] = {\n'% w1.shape[0]*w1.shape[1])
t = str(w1[1,0])
print(t)
for j in range(w1.shape[0]):
    for i in range(w1.shape[1]):
        f.write(str(w1[j,i])+',\n')
        pass
    pass
f.write('}\n')
f.close()
# plt.figure(1, figsize=(10, 3))
# plt.subplot(131)
# plt.title('net1')
# plt.scatter(x.data.numpy(), y.data.numpy())
# plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
#
# plt.subplot(132)
# plt.title('net2')
# plt.scatter(x.data.numpy(), y.data.numpy())
# plt.plot(x.data.numpy(), net2(x).data.numpy(), 'r-', lw=5)
#
# plt.subplot(133)
# plt.title('net3')
# plt.scatter(x.data.numpy(), y.data.numpy())
# plt.plot(x.data.numpy(), net3(x).data.numpy(), 'r-', lw=5)
# plt.show()