Pytorch中torch.nn使用

本文將介紹：

• torch.nn包
• 定義一個簡單的nn架構
• 定義優化器、損失函數
• 梯度的反向傳播

將使用LeNet-5架構進行說明

 

 

一、torch.nn包

torch.nn包來構建網絡；

torch.nn.Module類作為自定義類的基類；

nn.Module，包含了所有神經網絡層，比如卷積層或者是線性層；

torch.nn.Functional包，可以定義在前向傳播的時候的運算；比如，卷積、dropout以及激活函數

 

二、定義NN

我們將使用上述的類和包來定義NN

import torch
import torch.nn as nn
import torch.nn.functional as F

接下來，我們將使用torch.nn.Module來構建我們的NN

請記住，我們將嘗試構造出Lenet-5架構，代碼如下：

代碼詳解：

line3 在__init__函數中，我們首先調用了super()函數。這確保了我們繼承了nn.Module中的所有方法。現在我們就可以使用nn.Module中所有的方法和層。

從line4開始，網絡開始構建。

原始數據輸入 1, 32*32

首先有一個Conv2d(1, 6, (5, 5))卷基層；Conv2d參數為（輸入的通道，輸出通道，（核*核））；6, 28*28

池化后：6，14*14

ConvConv2d(6, 16, (5, 5))：輸出： 16，10*10

池化后：16， 5*5

拉平：16*5*5=400

最后經過三個線性層，輸出10個維度的特征

需要注意的是：我們旨在__init__()中定義層，所有的操作都是在forward中進行的。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, (5, 5)) # (input image channel, output channels, kernel size)
        self.conv2 = nn.Conv2d(6, 16, (5, 5))
        # in linear layer (output channels from conv2d x width x height)
        self.fc1 = nn.Linear(16 * 5 * 5, 120) # (in_features, out_features)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
        x = x.view(x.size(0), -1) ##這里卷積完成后，需要對其展平
        x = F.relu(self.fc1(x)) #調用激活函數
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
model = Net()  #實例化一個NN類
print(model)

 

output：

Net(
  (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)

 

三、定義優化器和損失函數

為了定義優化器，我們需要輸入torch.optim，

優化器具有很多選擇，如RMSprop， Adam, SGD, Adadelta等

# loss function and optimizer
import torch.optim as optim
loss_function = nn.MSELoss()
optimizer = optim.RMSprop(model.parameters(), lr=0.001)

optmizer接收兩個參數，第一個是我們定義的模型參數，第二個是學習率。

 

四、輸入的數據和反向傳播

在此，我們生成隨機數據對網絡進行實現。

input = torch.rand(1, 1, 32, 32)
out = model(input)
print(out, out.shape)

輸出：

tensor([[-0.0265, -0.0181,  0.1301,  0.0465,  0.0697,  0.0765, -0.0022,  0.0215,
           0.0908, -0.1489]], grad_fn=AddmmBackward) torch.Size([1, 10])

 

同樣的，如果我們需要定義一個真實的目標（label），因為我們需要計算損失。

在我們定義好我們的真實目標數據后，我們需要將其reshape成為[1,10]。

這是因為必須和輸出的形狀是相同的才可以計算損失。

# dummy targets
labels = torch.rand(10)
labels = labels.view(1, -1) # resize for the same shape as output
print(labels, labels.shape)

輸出：

tensor([[0.7737, 0.7730, 0.1154, 0.4876, 0.5071, 0.3506, 0.3078, 0.4576, 0.0926,
         0.1268]]) torch.Size([1, 10])

 

五、反向傳播

loss = loss_function(out, labels)
loss.backward()
print(loss)

輸出：

tensor(0.2090, grad_fn=MseLossBackward)

 

需要注意的是，我們現在只有一個輸入。

實際上，我們將在train循環使用上面的代碼。每次循環的時候需要將梯度設為0.

梯度設為0的操作，可以通過optimizer.zero_grad()來實現；