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()