PyTorch中如何将nn.Parameter替换为自定义nn.Module？

我想用自定义的nn.Module层替换模型内部的nn.Parameter参数，以下是简化后的示例代码：

import torch
import torch.nn as nn

class change_to_layer(nn.Module):
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.randn(100, 100))
    
    def __mul__(self, other):
        return self.forward(other)
    
    def __rmul__(self, other):
        return self.forward(other)

    def forward(self, x):
        return x @ self.w


class simple_model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(100, 100)
        self.scale = nn.Parameter(torch.ones(1))
        self.fc2 = nn.Linear(100, 100)

    def forward(self, x):
        x = self.fc1(x)
        x = self.scale * x
        x = self.fc2(x)
        print(x)


model = simple_model()

model.scale = change_to_layer()  # 尝试将nn.Parameter替换为nn.Module，触发错误

input = torch.randn(100)
print(model(input))

运行代码时触发如下错误：

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-15-1a4295f999b1> in <cell line: 33>()
     31 model = simple_model()
     32 
---> 33 model.scale = change_to_layer()
     34 
     35 input = torch.randn(100)

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in __setattr__(self, name, value)
   1633         elif params is not None and name in params:
   1634             if value is not None:
-> 1635                 raise TypeError("cannot assign '{}' as parameter '{}' "
   1636                                 "(torch.nn.Parameter or None expected)"
   1637                                 .format(torch.typename(value), name))

TypeError: cannot assign '__main__.change_to_layer' as parameter 'scale' (torch.nn.Parameter or None expected)

请问如何修改该类变量的类型以实现替换？

方法1：修改模型初始化逻辑（推荐） #

直接在模型初始化时根据需求选择用nn.Parameter还是自定义层，避免后续动态替换的麻烦。可以加个开关参数控制：

import torch
import torch.nn as nn

class change_to_layer(nn.Module):
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.randn(100, 100))
    
    def __mul__(self, other):
        return self.forward(other)
    
    def __rmul__(self, other):
        return self.forward(other)

    def forward(self, x):
        return x @ self.w


class simple_model(nn.Module):
    def __init__(self, use_custom_layer=False):
        super().__init__()
        self.fc1 = nn.Linear(100, 100)
        # 根据参数选择初始化类型
        if use_custom_layer:
            self.scale = change_to_layer()
        else:
            self.scale = nn.Parameter(torch.ones(1))
        self.fc2 = nn.Linear(100, 100)

    def forward(self, x):
        x = self.fc1(x)
        x = self.scale * x
        x = self.fc2(x)
        return x


# 用自定义层初始化模型
model = simple_model(use_custom_layer=True)
input = torch.randn(100)
print(model(input))

方法2：动态替换时手动调整注册表 #

PyTorch会把nn.Parameter存在_parameters字典，nn.Module存在_modules字典。scale已经被注册为参数，直接赋值模块会报错，所以要先移除参数注册，再添加模块注册：

import torch
import torch.nn as nn

class change_to_layer(nn.Module):
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.randn(100, 100))
    
    def __mul__(self, other):
        return self.forward(other)
    
    def __rmul__(self, other):
        return self.forward(other)

    def forward(self, x):
        return x @ self.w


class simple_model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(100, 100)
        self.scale = nn.Parameter(torch.ones(1))
        self.fc2 = nn.Linear(100, 100)

    def forward(self, x):
        x = self.fc1(x)
        x = self.scale * x
        x = self.fc2(x)
        return x


model = simple_model()

# 先从参数字典中移除scale
del model._parameters['scale']
# 将自定义层添加到模块字典中
model._modules['scale'] = change_to_layer()

input = torch.randn(100)
print(model(input))

补充说明 #

PyTorch的nn.Module通过__setattr__自动管理参数和子模块：赋值nn.Parameter时会加入_parameters，赋值nn.Module时加入_modules。一旦某个名字被注册为参数，后续再赋值非参数类型就会触发类型错误，所以动态替换必须手动调整这两个字典。

内容的提问来源于stack exchange，提问作者core_not_dumped