在许多地方都有修改原模型的需求，比如最简单和常见的例子就是修改模型最后的全链接层，来获得针对特定任务的输出结果。对于 PyTorch，可以通过多种方式来达到魔改模型的效果。

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

通过继承修改模型

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

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


class CustomizedResNet(ResNet):
    def __init__(self, block, 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):
        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


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)

通过组合修改模型

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

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

from torchvision.models import resnet18


class CustomizedResNet(torch.nn.Module):

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

    def forward(self, x):
        ...


my_resnet18 = CustomizedResNet(resnet18)

同样我们来实现跟上面同样的需求：

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())

通过猴子补丁修改模型

最后再介绍一个最简单粗暴的方法：猴子补丁（Monkey Patch）。之所以叫猴子补丁，是因为这种方法在程序设计的角度上看，是具有破坏性的。我们尽可能得还是应该使用上面两种方法。

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

from torchvision.models import resnet50

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

此外这种方法也仅能实现一些简单的需求，对于复杂的需求还是推荐使用组合的方法来完成。