在现有 PyTorch 层中实现自定义功能

问：

在 Pytorch 中，我想添加一个函数，
该函数对继承自 的所有层执行特定函数。nn.Module

在这里，图层表示：
///等。nn.Linearnn.Conv2dnn.MaxPool2d

首先，创建一个继承自 的新主类和层类。nn.Module

class Main(nn.Module):
    def foo(self, x):
        x += 2
        return x 

class Linear(nn.Module, Main):
    pass

class Conv2d(nn.Module, Main):
    pass

...

使用如上所述的多重继承，您可以将方法“foo”与现有的层功能一起使用。

我们将现有图层（如 A）称为 A，将新创建的图层类称为 B。nn.Linear

在组成新模型时，使用 B 时没有问题。
例如）

class NewModel(nn.Module):
    def __init__(self):
        super(NewModel, self).__init__()
        self.linear = Linear(64, 64)  # this Linear object is B, 'foo' method available

但是，使用 A 的模型（torchvision 中的内部类 Cifar10、ResNet50 等）
不能使用“foo”方法。

我还考虑过将“foo”方法放入使用 A 的对象中，并具有 ，setattr

必须打开 A 对象的所有模块并检查它们会带来不便，并且也很难跟踪内部计算。

在不修改库的情况下，有没有合适的方法可以解决上述问题？