PyTorch如何修改模型（魔改）

对模型缝缝补补、修修改改，是我们必须要掌握的技能，本文详细介绍了如何修改PyTorch模型？也就是我们经常说的如何魔改。👍

PyTorch 的模型是一个 torch.nn.Module 的某个子类的对象，修改模型实际就等价于修改某个类，对面向对象熟悉的同学应该知道，对类做修改有两个经典的方法：组合和继承。

1.修改模型层(模型框架⭐)

1.1通过继承修改模型

首先创建自己需要的模型类，然后其父类指向需要被修改的模型，这时自己的模型则具有完备的父类行为，最后在子类中实现魔改的逻辑。其大致的框架代码如下所示：

from torchvision.models import ResNet

class CustomizedResNet(ResNet):

    def __init__(self):
        super().__init__()
        ...
        
    def forward(self, x):
        ...

下面这个例子，将对 ResNet 进行魔改，把 ResNet 的 4 个 stage 输出的特征连接起来，然后通过一个全连接层后输出一个标量。

from torchvision.models.resnet import Bottleneck, BasicBlock, ResNet
import torch

# 定义一个自定义的ResNet类，继承自torchvision的ResNet类
class CustomizedResNet(ResNet):
    def __init__(self, block, layers, num_classes=2):
        """
        初始化函数
        block: ResNet中的基本块类型，可以是BasicBlock或Bottleneck
        layers: 每个层级的基本块数量，是一个列表
        num_classes: 输出的类别数量，默认为2
        """
        # 调用父类的初始化方法
        super().__init__(block, layers, num_classes)
        # 重新定义全连接层，改变输出的特征数量
        self.fc = torch.nn.Linear(int(512 * block.expansion * 1.875), num_classes)

    def forward(self, x):
        # 以下是ResNet的前向传播过程
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        # 通过四个残差层
        x1 = self.layer1(x)
        x2 = self.layer2(x1)
        x3 = self.layer3(x2)
        x4 = self.layer4(x3)

        # 将四个残差层的输出进行拼接
        x = torch.cat(
            [self.avgpool(x1),
             self.avgpool(x2),
             self.avgpool(x3),
             self.avgpool(x4),], dim=1)

        # 将拼接后的张量展平
        x = torch.flatten(x, 1)
        # 通过全连接层，得到最终的输出
        x = self.fc(x)

        return x

# 创建不同版本的ResNet模型
new_resnet34 = CustomizedResNet(BasicBlock, [3, 4, 6, 3], num_classes=1)
new_resnet50 = CustomizedResNet(Bottleneck, [3, 4, 6, 3], num_classes=1)
new_resnet101 = CustomizedResNet(Bottleneck, [3, 4, 23, 3], num_classes=1)
new_resnet200 = CustomizedResNet(Bottleneck, [3, 24, 36, 3], num_classes=1)

1.2通过组合修改模型(重点学👀)

在面向对象编程中，可能听说过「组合优于继承」，在模型修改的场景中其实也是这样，大多数情况下我们可能都适用组合而非继承。

首先依然需要创建模型的类，但这个类不再继承自魔改的类，而是直接继承 PyTorch 的模型基类 torch.nn.Module，然后将需要魔改的类作为类变量融入到模型中，下面是大致的框架代码：

from torchvision.models import resnet18
import torch.nn as nn

class CustomizedResNet(nn.Module):

    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
        ...

    def forward(self, x):
        ...

my_resnet18 = CustomizedResNet(resnet18)

同样，实现对 ResNet 进行魔改，把 ResNet 的 4 个 stage 输出的特征连接起来，然后通过一个全连接层后输出一个标量。

from torchvision.models import resnet50

class CustomizedResNet(torch.nn.Module):
    def __init__(self, backbone, num_classes=2):
        super().__init__()
        self.backbone = backbone
        self.fc = torch.nn.Linear(3840, num_classes)

    def forward(self, x):
        x = self.backbone.conv1(x)
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)

        x1 = self.backbone.layer1(x)
        x2 = self.backbone.layer2(x1)
        x3 = self.backbone.layer3(x2)
        x4 = self.backbone.layer4(x3)

        x = torch.cat(
            [
                self.backbone.avgpool(x1),
                self.backbone.avgpool(x2),
                self.backbone.avgpool(x3),
                self.backbone.avgpool(x4),
            ],
            dim=1,
        )

        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

new_resnet50 = CustomizedResNet(resnet50())

1.3通过猴子补丁修改模型

最简单粗暴的方法：猴子补丁（Monkey Patch）。之所以叫猴子补丁，是因为这种方法从程序设计的角度上来说，是具有破坏性的。而且这种方法仅能实现一些简单的修改需求，所以还是推荐使用继承或组合去修改我们的模型。😉

猴子补丁修改模型非常简单粗暴，直接使用需要修改的模型创建对象，然后直接对对象的属性做出修改。下面是把 ResNet34 的输出从 1000 改为 1 的简单例子：

from torchvision.models import resnet50
import torch.nn as nn

model = resnet50()
model.fc = nn.Linear(2048, 1)

还有一个例子，以 PyTorch 官方视觉库 torchvision 预定义好的模型 ResNet50 为例，修改模型的某一层或者某几层。先观察一下它的网络结构：

import torch
import torch.nn as nn
from collections import OrderedDict
import torchvision.models as models
net = models.resnet50()
print(net)

假设要用这个模型去做一个10分类的问题，就应该修改模型的 fc 层，将其输出节点数替换为10。另外，想再加一层全连接层。可以做如下修改：

classifier = nn.Sequential(OrderedDict([('fc1', nn.Linear(2048, 128)),
                          ('relu1', nn.ReLU()), 
                          ('dropout1',nn.Dropout(0.5)),
                          ('fc2', nn.Linear(128, 10)),
                          ('output', nn.Softmax(dim=1))
                          ]))

net.fc = classifier

这里的操作相当于将模型（net）最后名称为“fc”的层替换成了名称为“classifier”的结构。

2.添加外部输入

有时候在模型训练中，除了已有模型的输入之外，还需要输入额外的信息。比如在CNN网络中，我们除了输入图像，还需要同时输入图像对应的其他信息，这时候就需要在已有的CNN网络中添加额外的输入变量。基本思路是：将原模型添加输入位置前的部分作为一个整体，同时在forward中定义好原模型不变的部分、添加输入和后续层之间的连接关系，从而完成模型的修改。

以 torchvision 的 resnet50 模型为基础，任务还是10分类任务。不同点在于，我们希望利用已有的模型结构，在倒数第二层增加一个额外的输入变量 add_variable 来辅助预测。具体实现如下：

class Model(nn.Module):
    def __init__(self, net):
        super().__init__()
        self.net = net
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc_add = nn.Linear(1001, 10, bias=True)
        self.output = nn.Softmax(dim=1)
        
    def forward(self, x, add_variable):
        x = self.net(x)
        x = torch.cat((self.dropout(self.relu(x)),
                       add_variable.unsqueeze(1)),1)
        x = self.fc_add(x)
        x = self.output(x)
        return x

这里的实现要点是通过torch.cat实现了tensor的拼接。torchvision 中的 resnet50 输出是一个1000维的 tensor，通过修改 forward 函数，先将 1000 维的 tensor 通过激活函数层和dropout层，再和外部输入变量"add_variable"拼接，最后通过全连接层映射到指定的输出维度 10。

另外这里对外部输入变量"add_variable"进行 unsqueeze 操作是为了和 net 输出的 tensor 保持维度一致，常用于 add_variable 是单一数值 (scalar) 的情况，此时 add_variable 的维度是 (batch_size, )，需要在第二维补充维数1，从而可以和 tensor 进行torch.cat操作。
unsqueeze与sequeeze语法说明

最后，对我们修改好的模型结构进行实例化，就可以使用了：

net = models.resnet50()
model = Model(net).cuda()

另外别忘了，训练中在输入数据的时候要给两个inputs：

outputs = model(inputs, add_var)

3.添加额外输出

有时候在模型训练中，除了模型最后的输出外，我们需要输出模型某一中间层的结果，以施加额外的监督，获得更好的中间层结果。基本的思路是修改模型定义中 forward 函数的 return 变量。

依然以 resnet50 做 10 分类任务为例，在已经定义好的模型结构上，同时输出 1000 维的倒数第二层和 10 维的最后一层结果。具体实现如下：

class Model(nn.Module):
    def __init__(self, net):
        super().__init__()
        self.net = net
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc1 = nn.Linear(1000, 10, bias=True)
        self.output = nn.Softmax(dim=1)
        
    def forward(self, x, add_variable):
        x1000 = self.net(x)
        x10 = self.dropout(self.relu(x1000))
        x10 = self.fc1(x10)
        x10 = self.output(x10)
        return x10, x1000

之后，对我们修改好的模型结构进行实例化，就可以使用了：

net = models.resnet50()
model = Model(net).cuda()

out10, out1000 = model(inputs, add_var)

参考

• Chenglu’s Log

• Pytorch修改预训练模型的方法汇总

😃😃😃