如何逐层遍历并改写任意预训练CNN拓扑以生成双路元模型？

要解决嵌套层结构（如VGG的classifier、ResNet的BasicBlock）的遍历与元模型构建问题，核心是递归遍历PyTorch模型的所有层级子模块，并针对不同层类型做针对性替换/复制。以下是可落地的实现方案：

核心思路 #

1. 递归遍历所有子模块：PyTorch模型是树形结构，需递归处理所有层级的子模块（包括自定义Block、Sequential等容器），避免遗漏嵌套层。
2. 卷积层替换：将原卷积层替换为groups=2的双倍输出卷积（输入通道为奇数时降级为groups=1），复制原权重实现双分支等价输出。
3. 非卷积层双分支复制：对ReLU、BN、Linear等层，创建双分支模块，将卷积输出的前半/后半通道分别送入两个复制后的层，处理后再拼接输出。
4. 自定义模块兼容：通过反射获取自定义模块的初始化参数，确保新模块能正确继承原结构。

代码实现 #

import torch
import torch.nn as nn
import copy
import inspect

def get_init_args(module):
    """获取自定义模块的初始化参数（假设参数名与模块属性名一致）"""
    sig = inspect.signature(module.__init__)
    args = {}
    for name, param in sig.parameters.items():
        if name == 'self':
            continue
        if hasattr(module, name):
            args[name] = getattr(module, name)
    return args

class DualBranchLayer(nn.Module):
    """通用双分支处理模块，适用于ReLU、BN、Linear等层"""
    def __init__(self, original_layer):
        super().__init__()
        self.layer1 = copy.deepcopy(original_layer)
        self.layer2 = copy.deepcopy(original_layer)
    
    def forward(self, x):
        split_size = x.size(1) // 2
        x1 = x[:, :split_size, ...]
        x2 = x[:, split_size:, ...]
        out1 = self.layer1(x1)
        out2 = self.layer2(x2)
        return torch.cat([out1, out2], dim=1)

def build_meta_module(original_module):
    """递归构建元模型"""
    # 处理容器模块：Sequential/ModuleList
    if isinstance(original_module, nn.Sequential):
        meta_layers = []
        for layer in original_module:
            meta_layers.append(build_meta_module(layer))
        return nn.Sequential(*meta_layers)
    elif isinstance(original_module, nn.ModuleList):
        meta_modules = []
        for mod in original_module:
            meta_modules.append(build_meta_module(mod))
        return nn.ModuleList(meta_modules)
    
    # 处理自定义嵌套模块（如ResNet的BasicBlock）
    elif len(list(original_module.children())) > 0:
        init_args = get_init_args(original_module)
        meta_module = type(original_module)(**init_args)
        # 递归替换子模块
        for name, child in original_module.named_children():
            setattr(meta_module, name, build_meta_module(child))
        return meta_module
    
    # 处理叶子模块
    else:
        # 替换卷积层为groups=2的双倍输出卷积
        if isinstance(original_module, nn.Conv2d):
            out_channels = original_module.out_channels * 2
            if original_module.in_channels % 2 != 0:
                # 输入通道为奇数时，改用groups=1避免报错
                new_conv = nn.Conv2d(
                    in_channels=original_module.in_channels,
                    out_channels=out_channels,
                    kernel_size=original_module.kernel_size,
                    stride=original_module.stride,
                    padding=original_module.padding,
                    dilation=original_module.dilation,
                    groups=1,
                    bias=original_module.bias is not None
                )
                # 复制权重和偏置
                new_conv.weight.data = torch.cat([original_module.weight.data, original_module.weight.data], dim=0)
                if original_module.bias is not None:
                    new_conv.bias.data = torch.cat([original_module.bias.data, original_module.bias.data], dim=0)
            else:
                # 输入通道为偶数，使用groups=2
                new_conv = nn.Conv2d(
                    in_channels=original_module.in_channels,
                    out_channels=out_channels,
                    kernel_size=original_module.kernel_size,
                    stride=original_module.stride,
                    padding=original_module.padding,
                    dilation=original_module.dilation,
                    groups=2,
                    bias=original_module.bias is not None
                )
                # 拆分原权重适配groups=2的结构
                weight_part1 = original_module.weight[:, :original_module.in_channels//2, :, :]
                weight_part2 = original_module.weight[:, original_module.in_channels//2:, :, :]
                new_conv.weight.data = torch.cat([weight_part1, weight_part2], dim=0)
                if original_module.bias is not None:
                    new_conv.bias.data = torch.cat([original_module.bias.data, original_module.bias.data], dim=0)
            return new_conv
        
        # 处理其他叶子层（ReLU、BN、Linear等）
        elif isinstance(original_module, (nn.ReLU, nn.ReLU6, nn.BatchNorm2d, nn.Linear, nn.MaxPool2d, nn.AvgPool2d)):
            return DualBranchLayer(original_module)
        
        # 未定义的层类型直接复制
        else:
            return copy.deepcopy(original_module)

使用示例 #

# 加载预训练VGG16并构建元模型
original_vgg = torch.hub.load('pytorch/vision:v0.10.0', 'vgg16', pretrained=True)
meta_vgg = build_meta_module(original_vgg)

# 测试前向传播
x = torch.randn(1, 3, 224, 224)
output = meta_vgg(x)
print(output.shape)  # 输出: torch.Size([1, 2000])（原VGG16输出1000类，元模型翻倍）

# 加载预训练ResNet18并构建元模型
original_resnet = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
meta_resnet = build_meta_module(original_resnet)
output_resnet = meta_resnet(x)
print(output_resnet.shape)  # 输出: torch.Size([1, 2000])

关键细节说明 #

• 递归遍历：通过named_children()递归访问所有层级子模块，确保自定义Block（如BasicBlock）内部的卷积、BN等层被正确处理。
• 卷积层权重复制：针对groups=2的要求，将原权重拆分为两个输入通道组，拼接后适配新卷积的权重形状；输入通道为奇数时自动降级为groups=1，避免运行时错误。
• 双分支数据流：DualBranchLayer自动将输入按通道拆分，分别送入两个复制后的层，处理后拼接输出，保证后续层仍能接收双倍通道的输入。

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