在PyTorch中nn.Module類是用於定義網路中前向結構的父類，當要定義自己的網路結構時就要繼承這個類。現有的那些類式介面（如nn.Linear、nn.BatchNorm2d、nn.Conv2d等）也是繼承這個類的，nn.Module類可以巢狀若干nn.Module的物件，來形成網路結構的巢狀組合，下面記錄nn.Module的功能。

1.繼承nn.Module類的模組

使用其初始化函式建立物件，然後呼叫forward函式就能使用裡面的前向計算過程。

包括：Linear、ReLU、Sigmoid、Conv2d、ConvTransposed2d、Dropout...

2.容器nn.Sequential()

nn.Sequential是一個Sequential容器，模組將按照建構函式中傳遞的順序新增到模組中。通俗的說，就是根據自己的需求，把不同的函式組合成一個（小的）模組使用或者把組合的模組新增到自己的網路中。

 1 conv_module = nn.Sequential(
 2           nn.Conv2d(1,20,5),
 3           nn.ReLU(),
 4           nn.Conv2d(20,64,5),
 5           nn.ReLU()
 6         )
 7  
 8 # 具體的使用方法
 9 class Net(nn.Module):
 
10     def __init__(self):
11         super(Net, self).__init__()
12         self.conv_module = nn.Sequential(
13           nn.Conv2d(1,20,5),
14           nn.ReLU(),
15           nn.Conv2d(20,64,5),
16           nn.ReLU()
17         )
18  
19     def forward(self, input):
20         out = self.conv_module(input)
 
21         return out

Tip：

• 使用nn.Module，我們可以根據自己的需求改變傳播過程，如RNN等;
• 如果需要快速構建或者不需要過多的過程，直接使用nn.Sequential。

3.模組內部引數管理

可以用.parameters()或者.named_parameters()返回 其內的所有引數的迭代器：

 1 from torch import nn
 2 
 3 net=nn.Sequential(
 4     nn.Linear(4,2),          #輸入維度4，輸出維度2的線性層
 5     nn.Linear(2,2)
 6     )
 7 
 8 print(list(net.parameters()))
 9 # [Parameter containing:
10 # tensor([[-0.0829,  0.3424,  0.4514, -0.3981],
11 #         [-0.3401,  0.1429, -0.4525,  0.4991]], requires_grad=True),
12 # Parameter containing:
13 # tensor([-0.0321,  0.0872], requires_grad=True), 
14 # Parameter containing:
15 # tensor([[ 0.0628,  0.3092],
16 #         [ 0.5135, -0.4738]], requires_grad=True), 
17 # Parameter containing:
18 # tensor([-0.4249,  0.3921], requires_grad=True)]
19 
20 print(dict(net.named_parameters()))
21 # {'0.weight': Parameter containing:
22 # tensor([[-0.0829,  0.3424,  0.4514, -0.3981],
23 #         [-0.3401,  0.1429, -0.4525,  0.4991]], requires_grad=True), 
24 # '0.bias': Parameter containing:
25 # tensor([-0.0321,  0.0872], requires_grad=True), 
26 # '1.weight': Parameter containing:
27 # tensor([[ 0.0628,  0.3092],
28 #         [ 0.5135, -0.4738]], requires_grad=True), 
29 # '1.bias': Parameter containing:
30 # tensor([-0.4249,  0.3921], requires_grad=True)}

Tip：

• 以第0層為例，weight.shape=[2,4]，輸出維度在前，輸入維度在後，和Linear定義的時候相反；
• 相比.parameters()，.named_parameters()能看到引數名，預設情況下會使用所在的層數+引數型別的方式，從0層開始編號；
• 使用優化器時，可以直接呼叫nn.Module類定義的引數。

1 optimizer = optim.SGD(net.parameters(), lr=learning_rate)

4.模組樹形結構

模組之間通過巢狀組合會形成樹形結構，使用.children()可以獲取其直接孩子結點，使用.modules()可以獲取其所有子結點。

 1 from torch import nn
 2 
 3 class BaseNet(nn.Module):
 4 
 5     def __init__(self):
 6         super(BaseNet, self).__init__()
 7         self.net = nn.Linear(4, 3)               #輸入4維輸出3維的線性層
 8     def forward(self, x):
 9         return self.net(x)
10 
11 
12 class MyNet(nn.Module):
13 
14     def __init__(self):
15         super(MyNet, self).__init__()
16         self.net = nn.Sequential(                #使用Seq容器組合了三個模組
17             BaseNet(),
18             nn.ReLU(),
19             nn.Linear(3, 2)
20         )
21     def forward(self, x):
22         return self.net(x)
23 
24 
25 my_net = MyNet()
26 
27 print(list(my_net.children()))  # 直接孩子
28 # [Sequential(
29 #   (0): BaseNet((net): Linear(in_features=4, out_features=3, bias=True))
30 #   (1): ReLU()
31 #   (2): Linear(in_features=3, out_features=2, bias=True)
32 # )]
33 
34 print(list(my_net.modules()))  # 所有孩子
35 # [MyNet(
36 #   (net): Sequential(
37 #     (0): BaseNet((net): Linear(in_features=4, out_features=3, bias=True))
38 #     (1): ReLU()
39 #     (2): Linear(in_features=3, out_features=2, bias=True)
40 #   )), 
41 # Sequential(
42 #   (0): BaseNet((net): Linear(in_features=4, out_features=3, bias=True))
43 #   (1): ReLU()
44 #   (2): Linear(in_features=3, out_features=2, bias=True)
45 # ), 
46 # BaseNet((net): Linear(in_features=4, out_features=3, bias=True)), 
47 # Linear(in_features=4, out_features=3, bias=True), 
48 # ReLU(), 
49 # Linear(in_features=3, out_features=2, bias=True)]

.children()只返回自己的直系孩子列表，在這裡也就是一個nn.Sequential容器。而使用.modules()獲取的所有孩子是包括自己的。

5.裝置

使用.to(device)可以在具體的CPU/GPU上切換，這會將其所有子模組也一起轉移過去執行。

1 device = torch.device('cuda')
2 net = Net()
3 net.to(device)

Tip：模組的.to(device)是原地操作並返回自己的引用，而Tensor的.to(device)不會在當前Tensor上操作，返回的才是在目標裝置上對應建立的Tensor，所以net = MLP().to(device)。

6.載入和儲存

使用torch.load()載入檢查點檔案，然後傳入net.load_state_dict()網路模型設定引數，把當前類所有狀態net.state_dict()傳入torch.save()儲存到檔案中去。在訓練過程中，每隔一定的迭代次數可以儲存一下檢查點，將當前網路模型的狀態傳進去。

eg.ckpt.mdl是網路的一箇中間狀態

1 net.load_state_dict(torch.load('ckpt.mdl'))
2 #train...
3 torch.save(net.state_dict(), 'ckpt.mdl')

7.訓練和測試模式

前面的學習中提到Dropout和Batch Normalization在訓練和測試中的行為不同，需要對每一個nn.Module()模組單獨設定訓練狀態和測試狀態，可以直接為網路使用方法.train()切換到訓練模式，使用.eval()方法切換到測試模式。

1 #train
2 net.train()
3 ...
4 #test
5 net.eval()
6 ...

8.實現自定義的類

8.1

例如一個將資料只保留第一維，其它維合併的Flatten層，可以用在卷積層和全連線層之間：

 1 class Flatten(nn.Module):
 2 
 3     def __init__(self):
 4         super(Flatten, self).__init__()
 5     def forward(self, input):
 6         return input.view(input.size(0), -1)
 7 
 8 
 9 class TestNet(nn.Module):
10 
11     def __init__(self):
12         super(TestNet, self).__init__()
13         self.net = nn.Sequential(nn.Conv2d(1, 16, stride=1, padding=1),
14                                  nn.MaxPool2d(2, 2),
15                                  Flatten(),
16                                  nn.Linear(1*14*14, 10))
17     def forward(self, x):
18         return self.net(x)

Tip：只有類才能寫到Sequential裡面，比如F.relu不可以，要重新定義nn.ReLU

8.2

如果在繼承nn.Module類來實現模組時，出現需要操作Tensor的部分，那麼應當使用nn.Parameters（注意這裡P大寫）將其包裝起來。如果直接使用Tensor，那麼就不能用.parameters()（注意這裡p小寫）獲取到所有引數，也就不能直接傳給優化器去記錄要優化的這些引數了。

 1 class MyLinear(nn.Module):
 2     def __init__(self, inp, outp):
 3         super(MyLinear, self).__init__()
 4         # 線性層的引數w和b,對w而言輸出維度放在前面
 5         self.w = nn.Parameter(torch.randn(outp, inp))
 6         self.b = nn.Parameter(torch.randn(outp))
 7 
 8     def forward(self, x):
 9         x = x @ self.w.t() + self.b
10         return x

Tip：使用nn.Parameter()包裝Tensor時，自動設定了requires_grad=True，即預設情況下認為它是反向傳播優化的引數。