如何安全克隆PyTorch模块？重复创建新模块是否效率更高？

结论先行 #

你不能直接用.clone()方法实现模块复制，这个方法是PyTorch张量的专属方法，无法作用于nn.Module子类的实例，直接调用会报错。如果想要降低重复初始化的开销，可以用深拷贝的方式替代原有写法，且只要做对参数处理，不会改变网络的训练/测试行为。

正确的实现方式 #

你可以先初始化一个基准模块，再通过copy.deepcopy生成其他相同的实例，示例代码如下：

import copy

MIM_N_cell = []
MIM_S_cell = []
# 仅初始化一次基准模块
base_cell = MIM_NS_cell(
    input_dim=self.hidden_dim,
    hidden_dim=self.hidden_dim,
    kernel_size=self.kernel_size,
    model_cfg=model_cfg
)
for _ in range(self.num_layers - 1):
    # 深拷贝得到完全独立的模块实例
    MIM_N_cell.append(copy.deepcopy(base_cell))
    MIM_S_cell.append(copy.deepcopy(base_cell))

self.MIM_N_cell = nn.ModuleList(MIM_N_cell)
self.MIM_S_cell = nn.ModuleList(MIM_S_cell)

效果和注意事项 #

• 性能：深拷贝的开销远低于多次调用模块构造函数，尤其是模块内部结构复杂、包含大量子层时，提速效果非常明显。
• 行为一致性：深拷贝得到的模块参数完全独立，反向传播时梯度不会互相干扰，和你原有每次新建模块的运行逻辑完全一致。
• 特殊情况处理：如果你的MIM_NS_cell在构造时会执行随机初始化逻辑，深拷贝得到的所有模块都会和基准模块的初始参数完全相同，和原有写法中每个模块独立随机初始化的行为不同。如果你需要每个模块的初始参数独立随机，要么在深拷贝后对每个新实例单独调用自定义的初始化方法重置参数，要么保留原有新建逻辑。
• 兼容性：只要你的MIM_NS_cell没有绑定不可序列化的外部对象（比如动态生成的函数指针、外部文件句柄等），deepcopy都可以正常执行，绝大多数由标准PyTorch子模块组成的自定义类都满足这个要求。

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