PyTorch构建模型网络结构的6种方式

PyTorch提供了多种方式来构建模型的网络结构，我尝试总结一下，有如下6种常见方式（可能还有我没注意到的，欢迎补充）。我们平时写代码并不一定需要掌握全部方式，但是多了解一些，对于阅读理解别人的代码显然是有帮助的。

1，继承nn.Module类

这是构建自定义模型最基础也是最常见的方法。通过继承torch.nn.Module类，并在子类中定义__init__方法来初始化模型的各个层，以及在forward方法中定义数据的前向传播路径。

python">import torch  
import torch.nn as nn  
import torch.nn.functional as F  # 继承nn.Module  
class SimpleCNN(nn.Module):  def __init__(self, num_classes=10):  super(SimpleCNN, self).__init__()  # 定义卷积层  self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1)  self.relu1 = nn.ReLU()  self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  self.conv2 = nn.Conv2d(32, 64, 3, padding=1)  self.relu2 = nn.ReLU()  self.pool2 = nn.MaxPool2d(2, 2)  # 定义全连接层  self.fc1 = nn.Linear(64 * 8 * 8, 128)  self.relu3 = nn.ReLU()  self.fc2 = nn.Linear(128, num_classes)  def forward(self, x):  x = self.conv1(x)  x = self.relu1(x)  x = self.pool1(x)  x = self.conv2(x)  x = self.relu2(x)  x = self.pool2(x)  # 展平特征图  x = x.view(-1, 64 * 8 * 8)  # 全连接层  x = self.fc1(x)  x = self.relu3(x)  x = self.fc2(x)  return x  # 实例化模型并打印结构  
model = SimpleCNN(num_classes=10)  
print(model)  # 假设有一个输入张量，测试模型  
input_tensor = torch.randn(1, 3, 32, 32)  
output = model(input_tensor)  
print(output.shape)  # 应该是[1, 10]，表示10个类别的输出

优点：

1）高度灵活：允许用户定义任意复杂的前向传播逻辑，并可以轻松地插入自定义的操作或层。

2）功能强大：通过继承nn.Module，用户可以充分利用PyTorch提供的各种功能，如参数管理、模型保存/加载、GPU加速等。

缺点：

1）在模型层数多结构复杂时，只使用nn.Module类来编写会显得凌乱，后期难以维护

2）相比nn.Sequential，代码量更多，需要定义一个类，并且手动编写前向传播部分

2，使用nn.Sequential

对于顺序连接的层，可以使用nn.Sequential来简化模型的构建。nn.Sequential接受一个层列表作为输入，并自动定义前向传播。

python">import torch  
import torch.nn as nn  
import torch.nn.functional as F  # 定义一个简单的CNN模型，使用nn.Sequential  
class SimpleCNN(nn.Sequential):  def __init__(self, num_classes=10):  super(SimpleCNN, self).__init__()  # 添加卷积层  self.add_module('conv1', nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1))  self.add_module('relu1', nn.ReLU())  self.add_module('pool1', nn.MaxPool2d(kernel_size=2, stride=2))  self.add_module('conv2', nn.Conv2d(32, 64, 3, padding=1))  self.add_module('relu2', nn.ReLU())  self.add_module('pool2', nn.MaxPool2d(2, 2))  # 添加全连接层，注意需要先flatten特征图  self.add_module('flatten', nn.Flatten())  self.add_module('fc1', nn.Linear(64 * 8 * 8, 128))  # 修正了这里的输入维度  self.add_module('relu3', nn.ReLU())  self.add_module('fc2', nn.Linear(128, num_classes))  # 实例化模型并打印结构  
model = SimpleCNN(num_classes=10)  
print(model)  # 假设有一个输入张量  
input_tensor = torch.randn(1, 3, 32, 32)  
output = model(input_tensor)  
print(output.shape)  # 应该是[1, 10]，表示10个类别的输出

注：其中add.module是nn.Sequential中的一个方法用于向 nn.Sequential 容器中添加一个模块（即一个层或一个子网络）。当你创建一个 nn.Sequential 实例时，你可以通过调用 self.add_module 方法来逐个添加你想要的层。这个方法接受两个参数：name和module，module是一个 nn.Module 的实例，表示要添加的层或子网络。add.module不仅可以用在初始化方法中，还可以动态添加网络结构。

优点：

1）简单直观，代码量少，易于维护

缺点：

1）灵活性不足：对于需要复杂前向传播逻辑或者非线性层次结构（如跳跃结构或者分支结构）的模型，只用nn.Sequential不方便

2）调试不便：因为所有层都被封装在Sequential中

3）自定义操作受限：不方便插入自定义的逻辑和操作

3，结合使用nn.Module和nn.Sequential

使用nn.Module构建网络，但其中每个block都用nn.Sequential构建的方式，实际上结合了nn.Sequential和nn.Module两者的特点

python">import torch  
import torch.nn as nn  
import torch.nn.functional as F  # 定义一个简单的CNN模型，结合使用nn.Module和nn.Sequential  
class SimpleCNN(nn.Module):  def __init__(self, num_classes=10):  super(SimpleCNN, self).__init__()  # 使用nn.Sequential定义卷积层部分  self.features = nn.Sequential(  nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1),  nn.ReLU(),  nn.MaxPool2d(kernel_size=2, stride=2),  nn.Conv2d(32, 64, 3, padding=1),  nn.ReLU(),  nn.MaxPool2d(2, 2)  )  # 定义全连接层部分  self.fc1 = nn.Linear(64 * 8 * 8, 128)  self.relu3 = nn.ReLU()  self.fc2 = nn.Linear(128, num_classes)  def forward(self, x):  # 通过卷积层部分  x = self.features(x)  # 展平特征图  x = x.view(-1, 64 * 8 * 8)  # 通过全连接层部分  x = self.fc1(x)  x = self.relu3(x)  x = self.fc2(x)  return x  # 实例化模型并打印结构  
model = SimpleCNN(num_classes=10)  
print(model)  # 假设有一个输入张量  
input_tensor = torch.randn(1, 3, 32, 32)  
output = model(input_tensor)  
print(output.shape)  # 应该是[1, 10]，表示10个类别的输出

对于案例里这种过于简单的模型，这种混合方式的优势还不明显。但是对于复杂模型，这种组合的实现方式相比起前两种更常见。

4，使用nn.ModuleDict

nn.ModuleDict 是 PyTorch 中的一个类，它继承自 nn.Module，用于存储模块（modules）的字典。与普通的 Python 字典不同，nn.ModuleDict 中的模块会被自动注册为参数，这样它们就可以被识别为模型的一部分，并且在调用 .parameters() 或 .to(device) 等方法时，这些模块中的参数也会被包含在内。

nn.ModuleDict 的用法很简单，你可以像使用普通字典一样使用它，但是键（key）必须是字符串，值（value）必须是 nn.Module 的实例。

python">import torch  
import torch.nn as nn  
import torch.nn.functional as F  
from collections import OrderedDict  # 定义一个简单的CNN模型，使用nn.OrderedDict来组织层  
class SimpleCNN(nn.Module):  def __init__(self, num_classes=10):  super(SimpleCNN, self).__init__()  # 使用nn.OrderedDict定义所有层#不需要注册  self.layers = nn.ModuleDict({  'conv1': nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1),  'relu1': nn.ReLU(),  'pool1': nn.MaxPool2d(kernel_size=2, stride=2),  'conv2': nn.Conv2d(32, 64, 3, padding=1),  'relu2': nn.ReLU(),  'pool2': nn.MaxPool2d(2, 2),  'flatten': nn.Identity(),  # 使用nn.Identity作为占位符，实际展平操作在forward中实现  'fc1': nn.Linear(64 * 8 * 8, 128),  'relu3': nn.ReLU(),  'fc2': nn.Linear(128, num_classes)  })  def forward(self, x):  # 顺序通过所有层  for name,layer in self.layers.items():  x = layer(x)  if name == 'pool2':  # 在池化后检查是否需要展平  x = x.view(-1, 64 * 8 * 8)  # 展平操作print(x.shape)return x  # 实例化模型并打印结构  
model = SimpleCNN(num_classes=10)  
print(model)  # 假设有一个输入张量  
input_tensor = torch.randn(1, 3, 32, 32)  
output = model(input_tensor)  
print(output.shape)  # 应该是[1, 10]，表示10个类别的输出

优点：

1）会自动注册其中的模块，但是这一点也可能成为缺点，主要看需求。如果调用的是pytoch自带的层，用自动注册更省事。

缺点：

1）会自动注册其中的模块同样可能成为一个确定，如果需要自己编写层，就需要手动注册。

2）出错后定位问题位置和调试相对更困难：前三种实现方式里都有层次化的结构，而用nn.ModuleDict来实现会缺乏这样的结构信息。另外nn.ModuleDict允许动态修改和删除模块，这会增加出错几率和调试难度。

5，使用nn.OrderedDict

当需要对层进行命名以便后续访问时，还可以使用collections.OrderedDict

python">import torch  
import torch.nn as nn  
import torch.nn.functional as F  class SimpleCNN(nn.Module):  def __init__(self, num_classes=10):  super(SimpleCNN, self).__init__()  # 使用nn.OrderedDict定义所有层的顺序  self.layers = nn.OrderedDict([  ('conv1', nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1)),  ('relu1', nn.ReLU()),  ('pool1', nn.MaxPool2d(kernel_size=2, stride=2)),  ('conv2', nn.Conv2d(32, 64, 3, padding=1)),  ('relu2', nn.ReLU()),  ('pool2', nn.MaxPool2d(2, 2)),  ('flatten', nn.Flatten()),  # 使用nn.Flatten进行展平操作  ('fc1', nn.Linear(64 * 8 * 8, 128)),  ('relu3', nn.ReLU()),  ('fc2', nn.Linear(128, num_classes))  ])  # 将layers中的模块注册到当前Module  for name, module in self.layers.items():  self.add_module(name, module)  def forward(self, x):  # 顺序通过所有层  for name, layer in self.layers.items():  x = layer(x)  if name == 'flatten':  # 在展平层后打印形状  print(x.shape)  return x  # 实例化模型并打印结构  
model = SimpleCNN()  
print(model)  # 假设有一个输入张量  
input_tensor = torch.randn(1, 3, 32, 32)  
output = model(input_tensor)  
print(output.shape)  # 应该是[1, 10]，表示10个类别的输出

在PyTorch中，即使不用nn.OrderedDict也可以给每一层命名。但使用OrderedDict来管理层的顺序和注册仍然有其优势，特别是在构建复杂模型或需要动态修改模型结构时。

另外，我们可以注意到nn.OrderedDict和nn.ModuleDict都是用字典的来存储模块，那么他们的区别是什么呢？

最主要的区别有两点：1）nn.OrderedDict是严格保持元素顺序的，而nn.Module在原本的实现里是不保持元素顺序的，但是在python3.7及之后的版本里已经改为保持元素的插入顺序。

2）nn.OrderedDict其实不是pytoch中特有的类，而是python中的类，所以它是不会自动注册模块的。需要手动注册。

优缺点：

相比nn.ModuleDict，nn.OrderedDict不会自动注册模块，这个特点视情况可能成为优点或缺点。

6，调用预训练好的模块拼装成网络模型

这种方式在深度学习领域非常常见，特别是在迁移学习和微调（fine-tuning）的场景中。

示例代码如下，我们取预训练的resnet18前两个block，和我们自己实现的分类头组合成一个全新的模型。

python">import torch  
import torch.nn as nn  
from torchvision import models  class MyPartialResNet18(nn.Module):  def __init__(self, num_classes=10):  super(MyPartialResNet18, self).__init__()  # 加载预训练的resnet18模型  self.resnet18 = models.resnet18(pretrained=True)  # 冻结整个resnet18的参数  for param in self.resnet18.parameters():  param.requires_grad = False  # 只保留前两个block  self.features = nn.Sequential(  self.resnet18.conv1,  self.resnet18.bn1,  self.resnet18.relu,  self.resnet18.maxpool,  self.resnet18.layer1, self.resnet18.layer2  )self.gmp=  nn.AdaptiveMaxPool2d((1, 1))self.flatten = nn.Flatten()# 自定义一个全连接层  self.fc = nn.Linear(self.resnet18.layer2[-1].conv2.out_channels, num_classes)  def forward(self, x):  x = self.features(x)  x = self.gmp(x)x = self.flatten(x)x = self.fc(x)  return x  # 创建模型实例  
model = MyPartialResNet18(num_classes=10)  # 打印模型结构  
print(model)  # 假设你有一个输入tensor x  
x = torch.randn(1, 3, 224, 224)  
# 输出模型的预测  
output = model(x)  
print(output.shape)  # 应该是[1, 10]，表示10个类别的预测

优点：可以快速实现，可以利用现有的预训练参数

缺点：缺乏灵活性

总之以上方法都有各自优缺点，各位按照自己的实际需求选择。