PyTorch快速搭建神經網路及其儲存提取方法詳解

本篇文章主要介紹了PyTorch快速搭建神經網路及其儲存提取方法詳解，現在分享給大家，也給大家做個參考。一起過來看看吧

有時候我們訓練了一個模型, 希望儲存它下次直接使用,不需要下次再花時間去訓練 ，本節我們來講解一下PyTorch快速搭建神經網路及其儲存提取方法詳解

一、PyTorch快速搭建神經網路方法

先看實驗代碼：

import torch import torch.nn.functional as F  # 方法1，通過定義一個Net類來建立神經網路 class Net(torch.nn.Module):   def __init__(self, n_feature, n_hidden, n_output):     super(Net, self).__init__()     self.hidden = torch.nn.Linear(n_feature, n_hidden)     self.predict = torch.nn.Linear(n_hidden, n_output)    def forward(self, x):     x = F.relu(self.hidden(x))     x = self.predict(x)     return x  net1 = Net(2, 10, 2) print('方法1：\n', net1)  # 方法2 通過torch.nn.Sequential快速建立神經網路結構 net2 = torch.nn.Sequential(   torch.nn.Linear(2, 10),   torch.nn.ReLU(),   torch.nn.Linear(10, 2),   ) print('方法2：\n', net2) # 經驗證，兩種方法構建的神經網路功能相同，結構細節稍有不同  ''''' 方法1：  Net (  (hidden): Linear (2 -> 10)  (predict): Linear (10 -> 2) ) 方法2：  Sequential (  (0): Linear (2 -> 10)  (1): ReLU ()  (2): Linear (10 -> 2) ) '''

先前學習了通過定義一個Net類來構建神經網路的方法，classNet中首先通過super函數繼承torch.nn.Module模組的構造方法，再通過添加屬性的方式搭建神經網路各層的結構資訊，在forward方法中完善神經網路各層之間的串連資訊，然後再通過定義Net類對象的方式完成對神經網路結構的構建。

構建神經網路的另一個方法，也可以說是快速構建方法，就是通過torch.nn.Sequential，直接完成對神經網路的建立。

兩種方法構建得到的神經網路結構完全相同，都可以通過print函數來列印輸出網路資訊，不過列印結果會有些許不同。

二、PyTorch的神經網路儲存和提取

在學習和研究深度學習的時候，當我們通過一定時間的訓練，得到了一個比較好的模型的時候，我們當然希望將這個模型及模型參數儲存下來，以備後用，所以神經網路的儲存和模型參數提取重載是很有必要的。

首先，我們需要在需要儲存網路結構及其模型參數的神經網路的定義、訓練部分之後通過torch.save()實現對網路結構和模型參數的儲存。有兩種儲存方式：一是儲存年整個神經網路的的結構資訊和模型參數資訊，save的對象是網路net；二是只儲存神經網路的訓練模型參數，save的對象是net.state_dict()，儲存結果都以.pkl檔案形式儲存。

對應上面兩種儲存方式，重載方式也有兩種。對應第一種完整網路結構資訊，重載的時候通過torch.load(‘.pkl')直接初始化新的神經網路對象即可。對應第二種只儲存模型參數資訊，需要首先搭建相同的神經網路結構，通過net.load_state_dict(torch.load('.pkl'))完成模型參數的重載。在網路比較大的時候，第一種方法會花費較多的時間。

代碼實現：

import torch from torch.autograd import Variable import matplotlib.pyplot as plt  torch.manual_seed(1) # 設定隨機數種子  # 建立資料 x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1) y = x.pow(2) + 0.2*torch.rand(x.size()) x, y = Variable(x, requires_grad=False), Variable(y, requires_grad=False)  # 將待儲存的神經網路定義在一個函數中 def save():   # 神經網路結構   net1 = torch.nn.Sequential(     torch.nn.Linear(1, 10),     torch.nn.ReLU(),     torch.nn.Linear(10, 1),     )   optimizer = torch.optim.SGD(net1.parameters(), lr=0.5)   loss_function = torch.nn.MSELoss()    # 訓練部分   for i in range(300):     prediction = net1(x)     loss = loss_function(prediction, y)     optimizer.zero_grad()     loss.backward()     optimizer.step()    # 繪圖部分   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)    # 儲存神經網路   torch.save(net1, '7-net.pkl')           # 儲存整個神經網路的結構和模型參數   torch.save(net1.state_dict(), '7-net_params.pkl') # 只儲存神經網路的模型參數  # 載入整個神經網路的結構及其模型參數 def reload_net():   net2 = torch.load('7-net.pkl')   prediction = net2(x)    plt.subplot(132)   plt.title('net2')   plt.scatter(x.data.numpy(), y.data.numpy())   plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)  # 只載入神經網路的模型參數，神經網路的結構需要與儲存的神經網路相同的結構 def reload_params():   # 首先搭建相同的神經網路結構   net3 = torch.nn.Sequential(     torch.nn.Linear(1, 10),     torch.nn.ReLU(),     torch.nn.Linear(10, 1),     )    # 載入神經網路的模型參數   net3.load_state_dict(torch.load('7-net_params.pkl'))   prediction = net3(x)    plt.subplot(133)   plt.title('net3')   plt.scatter(x.data.numpy(), y.data.numpy())   plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)  # 運行測試 save() reload_net() reload_params()

實驗結果：