文章目录
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)
参考
😃😃😃