nn 包提供通用深度學習網絡的模塊集合，接收輸入張量，計算輸出張量，并保存權重。通常使用兩種途徑搭建 PyTorch 中的模型：nn.Sequential和 nn.Module。

nn.Sequential通過線性層有序組合搭建模型；nn.Module通過__init__ 函數指定層，然后通過?forward 函數將層應用于輸入，更靈活地構建自定義模型。

目錄

搭建線性層

通過nn.Sequential搭建

通過nn.Module搭建

獲取模型摘要

---

搭建線性層

使用 nn 包搭建線性層。線性層接收 64*1000 維的輸入，保存 1000*100 維的權重，并計算 64*100 維的輸出。

import torch
from torch import nn
input_tensor = torch.randn(64, 1000)
linear_layer = nn.Linear(1000, 100)
output = linear_layer(input_tensor)
print(input_tensor.size())
print(output.size())

通過nn.Sequential搭建

考慮一個兩層的神經網絡，四個節點作為輸入，五個節點在隱藏層，一個節點作為輸出

from torch import nn
model = nn.Sequential(nn.Linear(4, 5),nn.ReLU(),nn.Linear(5, 1),
)
print(model)

通過nn.Module搭建

在 PyTorch 中搭建模型的另一種方法是對 nn.Module?類進行子類化，通過__init__ 函數指定層，然后通過?forward 函數將層應用于輸入，更靈活地構建自定義模型。

考慮兩個卷積層和兩個完全連接層搭建的模型：

import torch.nn.functional as F
class Net(nn.Module):def __init__(self):super(Net, self).__init__()def forward(self, x):pass

定義__init__ 函數和forward 函數

def __init__(self):super(Net, self).__init__()self.conv1 = nn.Conv2d(1, 20, 5, 1)self.conv2 = nn.Conv2d(20, 50, 5, 1)self.fc1 = nn.Linear(4*4*50, 500)self.fc2 = nn.Linear(500, 10)
def forward(self, x):x = F.relu(self.conv1(x))x = F.max_pool2d(x, 2, 2)x = F.relu(self.conv2(x))x = F.max_pool2d(x, 2, 2) x = x.view(-1, 4*4*50)x = F.relu(self.fc1(x))x = self.fc2(x)return F.log_softmax(x, dim=1)

重寫兩個類函數并打印模型

重寫：子類中實現一個與父類的成員函數原型完全相同的函數

Net.__init__ = __init__
Net.forward = forward
model = Net()
print(model)

?查看模型位置

print(next(model.parameters()).device)

?

將模型移動至CUDA設備?

device = torch.device("cuda:0")
model.to(device)
print(next(model.parameters()).device)

獲取模型摘要

借助torchsummary包查獲取模型摘要

pip install torchsummary

from torchsummary import summary
summary(model, input_size=(1, 28, 28))