Pytorch: torch.nn

import torch as t
from torch import nn

class Linear(nn.Module): # 繼承nn.Module
    def __init__(self, in_features, out_features):
        super(Linear, self).__init__() # 等價於nn.Module.__init__(self)
        self.w = nn.Parameter(t.randn(in_features, out_features))
        self.b = nn.Parameter(t.randn(out_features))
    
    def forward(self, x):
        x = x.mm(self.w) # x.@(self.w)
        return x + self.b.expand_as(x)

• 自定義層Linear必須繼承nn.Module，並且在其構造函數中需調用nn.Module的構造函數，即super(Linear, self).__init__() 或nn.Module.__init__(self)，推薦使用第一種用法，盡管第二種寫法更直觀。
• 在構造函數__init__中必須自己定義可學習的參數，並封裝成Parameter，如在本例中我們把w和b封裝成parameter。parameter是一種特殊的Tensor，但其默認需要求導（requires_grad = True），感興趣的讀者可以通過nn.Parameter??，查看Parameter類的源代碼。
• forward函數實現前向傳播過程，其輸入可以是一個或多個tensor。
• 無需寫反向傳播函數，nn.Module能夠利用autograd自動實現反向傳播，這點比Function簡單許多。
• 使用時，直觀上可將layer看成數學概念中的函數，調用layer(input)即可得到input對應的結果。它等價於layers.__call__(input)，在__call__函數中，主要調用的是 layer.forward(x)，另外還對鈎子做了一些處理。所以在實際使用中應盡量使用layer(x)而不是使用layer.forward(x)，關於鈎子技術將在下文講解。
• Module中的可學習參數可以通過named_parameters()或者parameters()返回迭代器，前者會給每個parameter都附上名字，使其更具有辨識度。

可見利用Module實現的全連接層，比利用Function實現的更為簡單，因其不再需要寫反向傳播函數。

module中parameter的命名規范：

• 對於類似self.param_name = nn.Parameter(t.randn(3, 4))，命名為param_name
• 對於子Module中的parameter，會其名字之前加上當前Module的名字。如對於self.sub_module = SubModel()，SubModel中有個parameter的名字叫做param_name，那么二者拼接而成的parameter name 就是sub_module.param_name。

ReLU函數有個inplace參數，如果設為True，它會把輸出直接覆蓋到輸入中，這樣可以節省內存/顯存。之所以可以覆蓋是因為在計算ReLU的反向傳播時，只需根據輸出就能夠推算出反向傳播的梯度。但是只有少數的autograd操作支持inplace操作（如tensor.sigmoid_()），除非你明確地知道自己在做什么，否則一般不要使用inplace操作。

 

在以上的例子中，基本上都是將每一層的輸出直接作為下一層的輸入，這種網絡稱為前饋傳播網絡（feedforward neural network）。對於此類網絡如果每次都寫復雜的forward函數會有些麻煩，在此就有兩種簡化方式，ModuleList和Sequential。其中Sequential是一個特殊的module，它包含幾個子Module，前向傳播時會將輸入一層接一層的傳遞下去。ModuleList也是一個特殊的module，可以包含幾個子module，可以像用list一樣使用它，但不能直接把輸入傳給ModuleList。下面舉例說明。

# Sequential的三種寫法
net1 = nn.Sequential()
net1.add_module('conv', nn.Conv2d(3, 3, 3))
net1.add_module('batchnorm', nn.BatchNorm2d(3))
net1.add_module('activation_layer', nn.ReLU())

net2 = nn.Sequential(
        nn.Conv2d(3, 3, 3),
        nn.BatchNorm2d(3),
        nn.ReLU()
        )

from collections import OrderedDict
net3= nn.Sequential(OrderedDict([
          ('conv1', nn.Conv2d(3, 3, 3)),
          ('bn1', nn.BatchNorm2d(3)),
          ('relu1', nn.ReLU())
        ]))
print('net1:', net1)
print('net2:', net2)
print('net3:', net3)


net1: Sequential(
  (conv): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
  (batchnorm): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (activation_layer): ReLU()
)
net2: Sequential(
  (0): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
  (1): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU()
)
net3: Sequential(
  (conv1): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
  (bn1): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu1): ReLU()
)

modellist = nn.ModuleList([nn.Linear(3,4), nn.ReLU(), nn.Linear(4,2)])
input = t.randn(1, 3)
for model in modellist:
    input = model(input)
# 下面會報錯,因為modellist沒有實現forward方法
# output = modelist(input)

# 看到這里，讀者可能會問，為何不直接使用Python中自帶的list，而非要多此一舉呢？這是因為ModuleList是Module的子類，當在Module中使用它的時候，就能自動識別為子module。

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.list = [nn.Linear(3, 4), nn.ReLU()]
        self.module_list = nn.ModuleList([nn.Conv2d(3, 3, 3), nn.ReLU()])
    def forward(self):
        pass
model = MyModule()
model

MyModule(
  (module_list): ModuleList(
    (0): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
    (1): ReLU()
  )
)

可見，list中的子module並不能被主module所識別，而ModuleList中的子module能夠被主module所識別。這意味着如果用list保存子module，將無法調整其參數，因其未加入到主module的參數中。

除ModuleList之外還有ParameterList，其是一個可以包含多個parameter的類list對象。在實際應用中，使用方式與ModuleList類似。如果在構造函數__init__中用到list、tuple、dict等對象時，一定要思考是否應該用ModuleList或ParameterList代替。

損失函數：

# batch_size=3，計算對應每個類別的分數（只有兩個類別）
score = t.randn(3, 2)
# 三個樣本分別屬於1，0，1類，label必須是LongTensor
label = t.Tensor([1, 0, 1]).long()

# loss與普通的layer無差異
criterion = nn.CrossEntropyLoss()
loss = criterion(score, label)
loss

 

優化器：

from torch import  optim
optimizer = optim.SGD(params=net.parameters(), lr=1)
optimizer.zero_grad() # 梯度清零，等價於net.zero_grad()

input = t.randn(1, 3, 32, 32)
output = net(input)
output.backward(output) # fake backward

optimizer.step() # 執行優化

# 為不同子網絡設置不同的學習率，在finetune中經常用到
# 如果對某個參數不指定學習率，就使用最外層的默認學習率
optimizer =optim.SGD([
                {'params': net.features.parameters()}, # 學習率為1e-5
                {'params': net.classifier.parameters(), 'lr': 1e-2}
            ], lr=1e-5)
optimizer

# 只為兩個全連接層設置較大的學習率，其余層的學習率較小
special_layers = nn.ModuleList([net.classifier[0], net.classifier[3]])
special_layers_params = list(map(id, special_layers.parameters()))
base_params = filter(lambda p: id(p) not in special_layers_params,
                     net.parameters())

optimizer = t.optim.SGD([
            {'params': base_params},
            {'params': special_layers.parameters(), 'lr': 0.01}
        ], lr=0.001 )
optimizer

# 方法1: 調整學習率，新建一個optimizer
old_lr = 0.1
optimizer1 =optim.SGD([
                {'params': net.features.parameters()},
                {'params': net.classifier.parameters(), 'lr': old_lr*0.1}
            ], lr=1e-5)
optimizer1

# 方法2: 調整學習率, 手動decay, 保存動量
for param_group in optimizer.param_groups:
    param_group['lr'] *= 0.1 # 學習率為之前的0.1倍
optimizer

 

nn中還有一個很常用的模塊：nn.functional，nn中的大多數layer，在functional中都有一個與之相對應的函數。nn.functional中的函數和nn.Module的主要區別在於，用nn.Module實現的layers是一個特殊的類，都是由class layer(nn.Module)定義，會自動提取可學習的參數。而nn.functional中的函數更像是純函數，由def function(input)定義。下面舉例說明functional的使用，並指出二者的不同之處。

此時讀者可能會問，應該什么時候使用nn.Module，什么時候使用nn.functional呢？答案很簡單，如果模型有可學習的參數，最好用nn.Module，否則既可以使用nn.functional也可以使用nn.Module，二者在性能上沒有太大差異，具體的使用取決於個人的喜好。如激活函數（ReLU、sigmoid、tanh），池化（MaxPool）等層由於沒有可學習參數，則可以使用對應的functional函數代替，而對於卷積、全連接等具有可學習參數的網絡建議使用nn.Module。下面舉例說明，如何在模型中搭配使用nn.Module和nn.functional。另外雖然dropout操作也沒有可學習操作，但建議還是使用nn.Dropout而不是nn.functional.dropout，因為dropout在訓練和測試兩個階段的行為有所差別，使用nn.Module對象能夠通過model.eval操作加以區分。

from torch.nn import functional as F
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.pool(F.relu(self.conv1(x)), 2)
        x = F.pool(F.relu(self.conv2(x)), 2)
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

對於不具備可學習參數的層（激活層、池化層等），將它們用函數代替，這樣則可以不用放置在構造函數__init__中。對於有可學習參數的模塊，也可以用functional來代替，只不過實現起來較為繁瑣，需要手動定義參數parameter，如前面實現自定義的全連接層，就可將weight和bias兩個參數單獨拿出來，在構造函數中初始化為parameter。

 

初始化策略：

在深度學習中參數的初始化十分重要，良好的初始化能讓模型更快收斂，並達到更高水平，而糟糕的初始化則可能使得模型迅速癱瘓。PyTorch中nn.Module的模塊參數都采取了較為合理的初始化策略，因此一般不用我們考慮，當然我們也可以用自定義初始化去代替系統的默認初始化。而當我們在使用Parameter時，自定義初始化則尤為重要，因t.Tensor()返回的是內存中的隨機數，很可能會有極大值，這在實際訓練網絡中會造成溢出或者梯度消失。PyTorch中nn.init模塊就是專門為初始化而設計，如果某種初始化策略nn.init不提供，用戶也可以自己直接初始化。

 

nn.Module深入分析

如果想要更深入地理解nn.Module，究其原理是很有必要的。首先來看看nn.Module基類的構造函數：

def __init__(self):
    self._parameters = OrderedDict()
    self._modules = OrderedDict()
    self._buffers = OrderedDict()
    self._backward_hooks = OrderedDict()
    self._forward_hooks = OrderedDict()
    self.training = True

其中每個屬性的解釋如下：

• _parameters：字典，保存用戶直接設置的parameter，self.param1 = nn.Parameter(t.randn(3, 3))會被檢測到，在字典中加入一個key為'param'，value為對應parameter的item。而self.submodule = nn.Linear(3, 4)中的parameter則不會存於此。
• _modules：子module，通過self.submodel = nn.Linear(3, 4)指定的子module會保存於此。
• _buffers：緩存。如batchnorm使用momentum機制，每次前向傳播需用到上一次前向傳播的結果。
• _backward_hooks與_forward_hooks：鈎子技術，用來提取中間變量，類似variable的hook。
• training：BatchNorm與Dropout層在訓練階段和測試階段中采取的策略不同，通過判斷training值來決定前向傳播策略。

上述幾個屬性中，_parameters、_modules和_buffers這三個字典中的鍵值，都可以通過self.key方式獲得，效果等價於self._parameters['key'].下面舉例說明。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 等價與self.register_parameter('param1' ,nn.Parameter(t.randn(3, 3)))
        self.param1 = nn.Parameter(t.rand(3, 3))
        self.submodel1 = nn.Linear(3, 4) 
    def forward(self, input):
        x = self.param1.mm(input)
        x = self.submodel1(x)
        return x
net = Net()
net

net._modules
OrderedDict([('submodel1', Linear(in_features=3, out_features=4, bias=True))])

net._parameters
OrderedDict([('param1', Parameter containing:
              tensor([[ 0.3398,  0.5239,  0.7981],
                      [ 0.7718,  0.0112,  0.8100],
                      [ 0.6397,  0.9743,  0.8300]]))])

nn.Module在實際使用中可能層層嵌套，一個module包含若干個子module，每一個子module又包含了更多的子module。為方便用戶訪問各個子module，nn.Module實現了很多方法，如函數children可以查看直接子module，函數module可以查看所有的子module（包括當前module）。與之相對應的還有函數named_childen和named_modules，其能夠在返回module列表的同時返回它們的名字。

 

對於batchnorm、dropout、instancenorm等在訓練和測試階段行為差距巨大的層，如果在測試時不將其training值設為True，則可能會有很大影響，這在實際使用中要千萬注意。雖然可通過直接設置training屬性，來將子module設為train和eval模式，但這種方式較為繁瑣，因如果一個模型具有多個dropout層，就需要為每個dropout層指定training屬性。更為推薦的做法是調用model.train()函數，它會將當前module及其子module中的所有training屬性都設為True，相應的，model.eval()函數會把training屬性都設為False。

在PyTorch中保存模型十分簡單，所有的Module對象都具有state_dict()函數，返回當前Module所有的狀態數據。將這些狀態數據保存后，下次使用模型時即可利用model.load_state_dict()函數將狀態加載進來。優化器（optimizer）也有類似的機制，不過一般並不需要保存優化器的運行狀態。

# 保存模型
t.save(net.state_dict(), 'net.pth')

# 加載已保存的模型
net2 = Net()
net2.load_state_dict(t.load('net.pth'))

將Module放在GPU上運行也十分簡單，只需兩步：

model = model.cuda()：將模型的所有參數轉存到GPU
input.cuda()：將輸入數據也放置到GPU上
至於如何在多個GPU上並行計算，PyTorch也提供了兩個函數，可實現簡單高效的並行GPU計算

nn.parallel.data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None)
class torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
可見二者的參數十分相似，通過device_ids參數可以指定在哪些GPU上進行優化，output_device指定輸出到哪個GPU上。唯一的不同就在於前者直接利用多GPU並行計算得出結果，而后者則返回一個新的module，能夠自動在多GPU上進行並行加速。

# method 1
new_net = nn.DataParallel(net, device_ids=[0, 1])
output = new_net(input)

# method 2
output = nn.parallel.data_parallel(new_net, input, device_ids=[0, 1])
DataParallel並行的方式，是將輸入一個batch的數據均分成多份，分別送到對應的GPU進行計算，各個GPU得到的梯度累加。與Module相關的所有數據也都會以淺復制的方式復制多份，在此需要注意，在module中屬性應該是只讀的。

 

nn和autograd的關系

nn.Module利用的也是autograd技術，其主要工作是實現前向傳播。在forward函數中，nn.Module對輸入的tensor進行的各種操作，本質上都是用到了autograd技術。這里需要對比autograd.Function和nn.Module之間的區別：

• autograd.Function利用了Tensor對autograd技術的擴展，為autograd實現了新的運算op，不僅要實現前向傳播還要手動實現反向傳播
• nn.Module利用了autograd技術，對nn的功能進行擴展，實現了深度學習中更多的層。只需實現前向傳播功能，autograd即會自動實現反向傳播
• nn.functional是一些autograd操作的集合，是經過封裝的函數

作為兩大類擴充PyTorch接口的方法，我們在實際使用中應該如何選擇呢？如果某一個操作，在autograd中尚未支持，那么只能實現Function接口對應的前向傳播和反向傳播。如果某些時候利用autograd接口比較復雜，則可以利用Function將多個操作聚合，實現優化，正如第三章所實現的Sigmoid一樣，比直接利用autograd低級別的操作要快。而如果只是想在深度學習中增加某一層，使用nn.Module進行封裝則更為簡單高效。