文章目录
PyTorch如何修改模型(魔改)
对模型缝缝补补、修修改改,是我们必须要掌握的技能,本文详细介绍了如何修改PyTorch模型?也就是我们经常说的如何魔改。👍
PyTorch 的模型是一个 torch.nn.Module
的某个子类的对象,修改模型实际就等价于修改某个类 ,对面向对象熟悉的同学应该知道,对类做修改有两个经典的方法:组合 和继承。
1.修改模型层(模型框架⭐)
1.1通过继承修改模型
首先创建自己需要的模型类,然后其父类指向需要被修改的模型,这时自己的模型则具有完备的父类行为,最后在子类中实现魔改的逻辑。其大致的框架代码如下所示:
python
from torchvision.models import ResNet
class CustomizedResNet(ResNet):
def __init__(self):
super().__init__()
...
def forward(self, x):
...
下面这个例子,将对 ResNet 进行魔改,把 ResNet 的 4 个 stage 输出的特征连接起来,然后通过一个全连接层后输出一个标量。
python
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
,然后将需要魔改的类作为类变量融入到模型中,下面是大致的框架代码:
python
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 输出的特征连接起来,然后通过一个全连接层后输出一个标量。
python
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 的简单例子:
python
from torchvision.models import resnet50
import torch.nn as nn
model = resnet50()
model.fc = nn.Linear(2048, 1)
还有一个例子,以 PyTorch 官方视觉库 torchvision 预定义好的模型 ResNet50 为例,修改模型的某一层或者某几层。先观察一下它的网络结构:
python
import torch
import torch.nn as nn
from collections import OrderedDict
import torchvision.models as models
net = models.resnet50()
print(net)
假设要用这个模型去做一个10分类的问题,就应该修改模型的 fc 层,将其输出节点数替换为10。另外,想再加一层全连接层。可以做如下修改:
python
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 来辅助预测。具体实现如下:
python
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语法说明
最后,对我们修改好的模型结构进行实例化,就可以使用了:
python
net = models.resnet50()
model = Model(net).cuda()
另外别忘了,训练中在输入数据的时候要给两个inputs:
python
outputs = model(inputs, add_var)
3.添加额外输出
有时候在模型训练中,除了模型最后的输出外,我们需要输出模型某一中间层的结果,以施加额外的监督,获得更好的中间层结果。基本的思路是修改模型定义中 forward 函数的 return 变量。
依然以 resnet50 做 10 分类任务为例,在已经定义好的模型结构上,同时输出 1000 维的倒数第二层和 10 维的最后一层结果。具体实现如下:
python
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
之后,对我们修改好的模型结构进行实例化,就可以使用了:
python
net = models.resnet50()
model = Model(net).cuda()
out10, out1000 = model(inputs, add_var)
参考
😃😃😃