PyTorch中两个神经网络模型的逐层线性组合实现方法

核心思路 #

由于M1和M2架构完全一致，我们可以先初始化一个与二者结构相同的模型M，再逐层遍历两个模型的可训练参数（weight和bias），按照lambda向量对应的系数完成线性组合，将结果赋值给M的对应层参数即可。

实现步骤与代码示例 #

假设你的模型是自定义的nn.Module子类（示例用线性层堆叠的模型），lambda_vec为长度等于模型层数的PyTorch张量或列表，具体实现如下：

import torch
import torch.nn as nn

# 示例模型结构（可替换为你的实际模型）
class MyModel(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super().__init__()
        layers = []
        prev_size = input_size
        for size in hidden_sizes:
            layers.append(nn.Linear(prev_size, size))
            layers.append(nn.ReLU())
            prev_size = size
        layers.append(nn.Linear(prev_size, output_size))
        self.layers = nn.Sequential(*layers)

# 已训练完成的模型M1、M2
M1 = MyModel(10, [20, 15], 5)
M2 = MyModel(10, [20, 15], 5)
# 注：此处假设模型已完成训练，参数已更新

# n维lambda向量（n为模型中带weight/bias的层数，示例中是3个Linear层）
lambda_vec = torch.tensor([0.3, 0.7, 0.5])  # 每个元素取值范围[0,1]

# 初始化组合模型M，复制M1的结构
M = MyModel(10, [20, 15], 5)

# 筛选出所有带参数的层（跳过ReLU这类无参数层）
m1_layers = [layer for layer in M1.layers if isinstance(layer, nn.Linear)]
m2_layers = [layer for layer in M2.layers if isinstance(layer, nn.Linear)]
m_layers = [layer for layer in M.layers if isinstance(layer, nn.Linear)]

# 逐层组合weight和bias
for i, (m1_layer, m2_layer, m_layer) in enumerate(zip(m1_layers, m2_layers, m_layers)):
    lam = lambda_vec[i]
    # 组合weight参数
    m_layer.weight.data = lam * m1_layer.weight.data + (1 - lam) * m2_layer.weight.data
    # 组合bias参数
    m_layer.bias.data = lam * m1_layer.bias.data + (1 - lam) * m2_layer.bias.data

关键注意事项 #

• 层顺序匹配：必须保证m1_layers、m2_layers、m_layers中的层顺序完全对应，同架构模型的带参数层顺序天然一致；如果模型是手动定义子模块（如self.fc1、self.fc2），可直接按属性名遍历匹配。
• 参数操作方式：直接操作参数的data属性可避免构建计算图（仅用于组合预训练参数场景）；若后续需要微调组合后的模型，可去掉.data，直接赋值张量运算结果以保留计算图。
• lambda向量长度：必须与模型中带weight/bias的层数一致，比如模型有3个Linear层，lambda_vec长度就需为3。

通用化实现（适配任意复杂结构模型） #

如果模型包含嵌套子模块，可通过参数名匹配的方式递归组合：

def combine_models(M1, M2, lambda_vec):
    # 初始化组合模型，复制M1的结构
    M = type(M1)(*M1.args, **M1.kwargs)  # 需确保模型__init__参数可通过args/kwargs获取
    # 若无法直接初始化，也可使用M = deepcopy(M1)

    # 获取两个模型的参数列表，按名称排序保证顺序一致
    m1_params = sorted(M1.named_parameters(), key=lambda x: x[0])
    m2_params = sorted(M2.named_parameters(), key=lambda x: x[0])
    m_params = sorted(M.named_parameters(), key=lambda x: x[0])

    n_layers = len(lambda_vec)
    # 验证参数数量与lambda向量匹配（每个层对应weight和bias两个参数）
    assert len(m1_params) == 2 * n_layers, "参数数量与lambda向量长度不匹配"

    for i in range(n_layers):
        # 提取第i层的weight参数
        _, m1_weight = m1_params[2*i]
        _, m2_weight = m2_params[2*i]
        _, m_weight = m_params[2*i]

        # 提取第i层的bias参数
        _, m1_bias = m1_params[2*i + 1]
        _, m2_bias = m2_params[2*i + 1]
        _, m_bias = m_params[2*i + 1]

        lam = lambda_vec[i]
        m_weight.data = lam * m1_weight.data + (1 - lam) * m2_weight.data
        m_bias.data = lam * m1_bias.data + (1 - lam) * m2_bias.data

    return M

# 调用示例
M = combine_models(M1, M2, lambda_vec)

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