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

蓝海渔夫

已于 2024-08-23 13:20:38 修改

阅读量1k

点赞数 14

分类专栏： pytorch 文章标签： pytorch 人工智能 python

于 2024-08-23 13:20:18 首次发布

本文链接：https://blog.csdn.net/wwimhere/article/details/141401901

版权

pytorch 专栏收录该内容

17 篇文章 1 订阅

订阅专栏

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

1，继承nn.Module类

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

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接受一个层列表作为输入，并自动定义前向传播。

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两者的特点

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 的实例。

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

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，和我们自己实现的分类头组合成一个全新的模型。

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个类别的预测

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

缺点：缺乏灵活性

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