1. 程式人生 > >從頭學pytorch(九):模型構造

從頭學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)
)

既然SequentialModuleList都可以進行列表化構造網路,那二者區別是什麼呢。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.

ModuleDict

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類來構造模型。
  • SequentialModuleListModuleDict類都繼承自Module類。
  • Sequential不同,ModuleListModuleDict並沒有定義一個完整的網路,它們只是將不同的模組存放在一起,需要自己定義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())