如何将多个无中间非线性激活的卷积层合并为单个PyTorch层？

当连续卷积层之间没有非线性激活（如ReLU、Sigmoid等）时，这些卷积操作可数学等价为一个单一卷积层。以下是原理说明及适配不同padding、stride的PyTorch实现方案。

核心原理 #

多层卷积的叠加（无非线性）等价于输入与一个「组合核」的卷积。我们可以通过单位脉冲模板输入推导这个组合核：将仅含单个1的张量输入到多层卷积中，最终输出就是等效的单一卷积核——因为单位脉冲经过卷积后的输出就是核本身，多层叠加后自然得到核的组合结果。

同时需处理padding和stride的累积效应：

• 合并后的总stride是所有层stride的乘积（每层stride对输出下采样，总下采样率为各层stride相乘）
• 合并后的padding需通过原多层卷积的输入输出尺寸关系反推，确保输出尺寸完全一致

PyTorch实现代码 #

import torch
import torch.nn as nn

def merge_conv_layers(conv_layers):
    """
    合并多个连续无非线性激活的Conv2d层为单个Conv2d层
    
    参数:
        conv_layers (list): 待合并的Conv2d层列表，按执行顺序排列
    
    返回:
        merged_conv (nn.Conv2d): 合并后的等效Conv2d层
    """
    # 校验输入层类型
    assert all(isinstance(layer, nn.Conv2d) for layer in conv_layers), "所有待合并层必须是Conv2d类型"
    
    # 获取输入输出通道数
    in_channels = conv_layers[0].in_channels
    out_channels = conv_layers[-1].out_channels
    
    # 计算累积stride（假设H/W方向stride一致）
    total_stride = 1
    for layer in conv_layers:
        total_stride *= layer.stride[0]
    
    # 计算等效核的最小尺寸，生成足够大的模板输入避免边界截断
    kernel_sizes = [layer.kernel_size[0] for layer in conv_layers]
    def calc_combined_kernel_size(kernel_sizes):
        size = kernel_sizes[0]
        for k in kernel_sizes[1:]:
            size += k - 1
        return size
    combined_kernel_size = calc_combined_kernel_size(kernel_sizes)
    max_padding = max(layer.padding[0] for layer in conv_layers)
    template_size = combined_kernel_size + 2 * max_padding * len(conv_layers)
    center = template_size // 2

    # 逐个输入通道生成脉冲，提取完整等效核
    merged_kernel = []
    for in_ch in range(in_channels):
        template = torch.zeros(1, in_channels, template_size, template_size)
        template[0, in_ch, center, center] = 1.0
        with torch.no_grad():
            x = template
            for layer in conv_layers:
                x = layer(x)
        merged_kernel.append(x.squeeze(0))
    merged_kernel = torch.stack(merged_kernel, dim=1)  # 最终形状: (out_channels, in_channels, h, w)
    
    # 反推合并后的padding，确保输出尺寸与原多层一致
    def calc_output_size(input_size, layers):
        h, w = input_size
        for layer in layers:
            kh, kw = layer.kernel_size
            sh, sw = layer.stride
            ph, pw = layer.padding
            h = (h + 2*ph - kh) // sh + 1
            w = (w + 2*pw - kw) // sw + 1
        return h, w
    
    test_input_size = (100, 100)
    original_output_size = calc_output_size(test_input_size, conv_layers)
    merged_kh, merged_kw = merged_kernel.shape[2], merged_kernel.shape[3]
    
    # 根据输出尺寸公式反推padding值
    ph = ((original_output_size[0] - 1)*total_stride + merged_kh - test_input_size[0]) // 2
    pw = ((original_output_size[1] - 1)*total_stride + merged_kw - test_input_size[1]) // 2
    
    # 创建合并后的卷积层
    merged_conv = nn.Conv2d(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=(merged_kh, merged_kw),
        stride=(total_stride, total_stride),
        padding=(ph, pw),
        bias=False
    )
    merged_conv.weight.data = merged_kernel
    
    # 合并bias（如果原层带bias）
    if any(layer.bias is not None for layer in conv_layers):
        merged_bias = torch.zeros(out_channels)
        x_bias = torch.zeros(in_channels)
        for layer in conv_layers:
            if layer.bias is not None:
                x_bias = torch.nn.functional.conv2d(
                    x_bias.view(1, in_channels, 1, 1),
                    layer.weight,
                    bias=layer.bias,
                    stride=layer.stride,
                    padding=layer.padding
                ).view(layer.out_channels)
        merged_conv.bias = nn.Parameter(x_bias)
    
    return merged_conv

# 测试示例
if __name__ == "__main__":
    # 创建两个测试卷积层
    conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1, bias=True)
    conv2 = nn.Conv2d(16, 8, kernel_size=3, stride=2, padding=1, bias=True)
    
    # 合并卷积层
    merged_conv = merge_conv_layers([conv1, conv2])
    
    # 生成随机输入
    input_tensor = torch.randn(2, 3, 64, 64)
    
    # 计算原多层与合并层的输出
    with torch.no_grad():
        original_output = conv2(conv1(input_tensor))
        merged_output = merged_conv(input_tensor)
    
    # 验证输出一致性
    print(f"输出误差最大值: {torch.max(torch.abs(original_output - merged_output))}")
    print(f"输出是否近似相等: {torch.allclose(original_output, merged_output, atol=1e-6)}")

注意事项 #

• 仅适用于无非线性激活、无池化的场景；若包含eval模式的BN，需先将BN参数融合到卷积核后再合并
• 模板输入尺寸需足够大，避免边界截断导致等效核不完整
• 当前代码假设各层H/W方向的kernel_size、stride、padding一致，若不同需修改对应逻辑
• 分组卷积（groups>1）需特殊处理，当前代码仅支持普通卷积

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