從頭學pytorch(九):模型構造

模型構造

nn.Module

nn.Module是pytorch中提供的一個類,是所有神經網絡模塊的基類.我們自定義的模塊要繼承這個基類.

import torch
from torch import nn

class MLP(nn.Module):
    # 聲明帶有模型參數的層，這里聲明了兩個全連接層
    def __init__(self, **kwargs):
        # 調用MLP父類Module的構造函數來進行必要的初始化。這樣在構造實例時還可以指定其他函數
        # 參數，如“模型參數的訪問、初始化和共享”一節將介紹的模型參數params
        super(MLP, self).__init__(**kwargs)
        self.hidden = nn.Linear(784, 256) # 隱藏層
        self.act = nn.ReLU()
        self.output = nn.Linear(256, 10)  # 輸出層
         

    # 定義模型的前向計算，即如何根據輸入x計算返回所需要的模型輸出
    def forward(self, x):
        a = self.act(self.hidden(x))
        return self.output(a)

X = torch.rand(2, 784)
net = MLP()
print(net)
net(X)

輸出如下:

MLP(
  (hidden): Linear(in_features=784, out_features=256, bias=True)
  (act): ReLU()
  (output): Linear(in_features=256, out_features=10, bias=True)
)

Module的子類

torch中還提供了一些其他的類,方便我們構造模型.這些類也都是繼承自nn.Module.

• Sequential
• ModuleList
• ModuleDict

這些類的定義都位於torch/nn/modules/container.py

Sequential

當模型的前向計算為簡單串聯各個層的計算時，Sequential類可以通過更加簡單的方式定義模型。這正是Sequential類的目的：它可以接收一個子模塊的有序字典（OrderedDict）或者一系列子模塊作為參數來逐一添加Module的實例，而模型的前向計算就是將這些實例按添加的順序逐一計算。

# Example of using Sequential
model = nn.Sequential(
            nn.Conv2d(1,20,5),
            nn.ReLU(),
            nn.Conv2d(20,64,5),
            nn.ReLU()
        )

# Example of using Sequential with OrderedDict
model = nn.Sequential(OrderedDict([
            ('conv1', nn.Conv2d(1,20,5)),
            ('relu1', nn.ReLU()),
            ('conv2', nn.Conv2d(20,64,5)),
            ('relu2', nn.ReLU())
        ]))

ModuleList

ModuleList接收一個子模塊的列表作為輸入，然后也可以類似List那樣進行append和extend操作:

net = nn.ModuleList([nn.Linear(784, 256), nn.ReLU()])
net.append(nn.Linear(256, 10)) # # 類似List的append操作
print(net[-1])  # 類似List的索引訪問
print(net)
# net(torch.zeros(1, 784)) # 會報NotImplementedError

輸出：

Linear(in_features=256, out_features=10, bias=True)
ModuleList(
  (0): Linear(in_features=784, out_features=256, bias=True)
  (1): ReLU()
  (2): Linear(in_features=256, out_features=10, bias=True)
)

既然Sequential和ModuleList都可以進行列表化構造網絡，那二者區別是什么呢。ModuleList僅僅是一個儲存各種模塊的列表，這些模塊之間沒有聯系也沒有順序（所以不用保證相鄰層的輸入輸出維度匹配），而且沒有實現forward功能需要自己實現，所以上面執行net(torch.zeros(1, 784))會報NotImplementedError；而Sequential內的模塊需要按照順序排列，要保證相鄰層的輸入輸出大小相匹配，內部forward功能已經實現。

ModuleList的出現只是讓網絡定義前向傳播時更加靈活，見下面官網的例子。

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])

    def forward(self, x):
        # ModuleList can act as an iterable, or be indexed using ints
        for i, l in enumerate(self.linears):
            x = self.linears[i // 2](x) + l(x)
        return x

這里注意nn.ModuleList傳入的是一個python list．

nn.ModuleList([nn.Linear(10, 10)])

不要寫成了

nn.ModuleList(nn.Linear(10, 10))

另外，ModuleList不同於一般的Python的list，加入到ModuleList里面的所有模塊的參數會被自動添加到整個網絡中，下面看一個例子對比一下。

class Module_ModuleList(nn.Module):
    def __init__(self):
        super(Module_ModuleList, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10)])
    
class Module_List(nn.Module):
    def __init__(self):
        super(Module_List, self).__init__()
        self.linears = [nn.Linear(10, 10)]

net1 = Module_ModuleList()
net2 = Module_List()

print("net1:")
for p in net1.parameters():
    print(p.size())

print("net2:")
for p in net2.parameters():
    print(p)

輸出：

net1:
torch.Size([10, 10])
torch.Size([10])
net2:

可以看到net2是沒有parameters的.net1是有parameters的.因為net1用的是nn.ModuleList而不是python list．

ＭoduleDict

ModuleDict接收一個子模塊的字典作為輸入, 然后也可以類似字典那樣進行添加訪問操作:

net = nn.ModuleDict({
    'linear': nn.Linear(784, 256),
    'act': nn.ReLU(),
})
net['output'] = nn.Linear(256, 10) # 添加
print(net['linear']) # 訪問
print(net.output)
print(net)
# net(torch.zeros(1, 784)) # 會報NotImplementedError

輸出：

Linear(in_features=784, out_features=256, bias=True)
Linear(in_features=256, out_features=10, bias=True)
ModuleDict(
  (act): ReLU()
  (linear): Linear(in_features=784, out_features=256, bias=True)
  (output): Linear(in_features=256, out_features=10, bias=True)
)

和ModuleList一樣，ModuleDict實例僅僅是存放了一些模塊的字典，並沒有定義forward函數需要自己定義。同樣，ModuleDict也與Python的Dict有所不同，ModuleDict里的所有模塊的參數會被自動添加到整個網絡中。

總結一下

• 可以通過繼承Module類來構造模型。
• Sequential、ModuleList、ModuleDict類都繼承自Module類。
• 與Sequential不同，ModuleList和ModuleDict並沒有定義一個完整的網絡，它們只是將不同的模塊存放在一起，需要自己定義forward函數。

構造復雜模型

上面介紹的Sequential使用簡單,但靈活性不足.通常我們還是自定義類,繼承nn.Module,去完成更復雜的模型定義和控制.

class FancyMLP(nn.Module):
    def __init__(self, **kwargs):
        super(FancyMLP, self).__init__(**kwargs)
        
        self.rand_weight = torch.rand((20, 20), requires_grad=False) # 不可訓練參數（常數參數）
        self.linear = nn.Linear(20, 20)

    def forward(self, x):
        x = self.linear(x)
        # 使用創建的常數參數，以及nn.functional中的relu函數和mm函數
        x = nn.functional.relu(torch.mm(x, self.rand_weight.data) + 1)
        
        # 復用全連接層。等價於兩個全連接層共享參數
        x = self.linear(x)
        # 控制流，這里我們需要調用item函數來返回標量進行比較
        while x.norm().item() > 1:
            x /= 2
        if x.norm().item() < 0.8:
            x *= 10
        return x.sum()

X = torch.rand(2, 20)
net = FancyMLP()
print(net)
print(net(X))

輸出

FancyMLP(
  (linear): Linear(in_features=20, out_features=20, bias=True)
)
tensor(2.0396, grad_fn=<SumBackward0>)

這里在print(net)時的輸出,是和__init__函數保持一致的,比如

class FancyMLP(nn.Module):
    def __init__(self, **kwargs):
        super(FancyMLP, self).__init__(**kwargs)
        
        self.rand_weight = torch.rand((20, 20), requires_grad=False) # 不可訓練參數（常數參數）
        self.linear = nn.Linear(20, 20)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.linear(x)
        # 使用創建的常數參數，以及nn.functional中的relu函數和mm函數
        x = nn.functional.relu(torch.mm(x, self.rand_weight.data) + 1)
        
        # 復用全連接層。等價於兩個全連接層共享參數
        x = self.linear(x)
        # 控制流，這里我們需要調用item函數來返回標量進行比較
        while x.norm().item() > 1:
            x /= 2
        if x.norm().item() < 0.8:
            x *= 10
        return x.sum()

X = torch.rand(2, 20)
net = FancyMLP()
print(net)
print(net(X))

輸出

FancyMLP(
  (linear): Linear(in_features=20, out_features=20, bias=True)
  (relu): ReLU()
)
tensor(7.5126, grad_fn=<SumBackward0>)

盡管在forward()里並沒有用到self.relu.

自定義的模型依然可以和Sequential一起使用.因為再復雜,它也還是繼承自nn.Module

net = nn.Sequential(nn.Linear(30, 20),FancyMLP())