PyTorch -- nn.Module 快速实践

• 网络层父类 nn.Module：Evrey Layer is nn.Module (nn.Linear, nn,Conv2d …)
• 具体的，我们在定义自已的网络时：需要继承 nn.Module，并重新实现
  • __init__ 方法: 一般放置网络中具有可学习参数的层（如全连接层、卷积层等）
    • 也可放置不具有可学习参数的层(如ReLU、dropout 等)；or 直接在 forward 方法 直接用 nn.functional 来代替
    • 除了基础模块，还可以用 nn.Sequential 来定义复合层
  • forward 方法：实现各个层之间的连接关系

经典代码示例：

import torch
 
class MyNet(torch.nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()  
        self.conv1 = torch.nn.Conv2d(3, 32, 3, 1, 1)
        self.relu1=torch.nn.ReLU()
        self.max_pooling1=torch.nn.MaxPool2d(2,1)
 
        self.conv2 = torch.nn.Conv2d(3, 32, 3, 1, 1)
        self.relu2=torch.nn.ReLU()
        self.max_pooling2=torch.nn.MaxPool2d(2,1)
 
        self.dense1 = torch.nn.Linear(32 * 3 * 3, 128)
        self.dense2 = torch.nn.Linear(128, 10)
 
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)  					## 写法一
        # x = torch.relu(x)  			    ## 写法二 (__init__ 中不需要定义 relf.relu1) 
        # x = torch.nn.functional.relu(x)  	## 写法三 (__init__ 中不需要定义 relf.relu1) 
        x = self.max_pooling1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.max_pooling2(x)
        x = self.dense1(x)
        x = self.dense2(x)
        return x
 
model = MyNet()  

print(model) 的结果如下所示：可参照结果对照理解

MyNet(
  (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (relu1): ReLU()
  (max_pooling1): MaxPool2d(kernel_size=2, stride=1, padding=0, dilation=1, ceil_mode=False)
  (conv2): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (relu2): ReLU()
  (max_pooling2): MaxPool2d(kernel_size=2, stride=1, padding=0, dilation=1, ceil_mode=False)
  (dense1): Linear(in_features=288, out_features=128, bias=True)
  (dense2): Linear(in_features=128, out_features=10, bias=True)
)

Module 类的常见方法：

• .children(), .named_children(): 返回模型的直接子模块，不包括嵌套的子模块

  • 适合快速查看模型的主要组成部分

  >>> list(model.children())[0]
  Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))  # 对应 `conv1`
  

• .modules(), .named_modules(): 递归遍历模型本身及其所有嵌套的子模块

  • 适合查看模型的完整结构。

  >>> list(model.modules())[0]  
  MyNet(  # 对应模型本身
    (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (relu1): ReLU()
    (max_pooling1): MaxPool2d(kernel_size=2, stride=1, padding=0, dilation=1, ceil_mode=False)
    (conv2): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (relu2): ReLU()
    (max_pooling2): MaxPool2d(kernel_size=2, stride=1, padding=0, dilation=1, ceil_mode=False)
    (dense1): Linear(in_features=288, out_features=128, bias=True)
    (dense2): Linear(in_features=128, out_features=10, bias=True)
  )
  

此外，使用 nn.Module 定义 model 可以使下列操作方便实现：

• 可转移到指定 device = torch.device('cuda)，使用 model.to(device)
• 可加载和保存 model 参数：
  • 加载: model.load_state_dict(torch.load('xxx.mdl'))
  • 保存: torch.save(model.state_dict(), 'xxx.del')
• 可进行状态转化 train/test：model.train() 和 model.eval() 一键切换

• B站视频参考资料
• 博客参考资料