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