如何高效合并两个基于通道范数的三维张量

在深度学习的卷积神经网络中，特征图通常表示为三维张量，形状为 C×H×W 或 H×W×C。通道范数指的是对每个通道单独计算一个标量值，常用的是 L2 范数，也就是该通道所有空间位置数值平方和再开根号。这个标量可以反映该通道整体激活强度，许多剪枝算法和特征选择方法都依赖它来判断通道的重要性。当需要把两个来自不同分支、不同层或不同模型的三维张量合并成一个张量时，直接按通道维度做拼接虽然简单，但会带来两个问题：一是通道数翻倍，增加后续计算负担；二是两个张量中可能存在大量低范数的冗余通道，直接拼接后这些冗余通道依然参与运算，浪费算力。因此，先计算每个通道的范数，再根据范数设计合并策略，是一种高效且可解释的做法。

本文会从通道范数的计算与意义出发，介绍三种基于范数的合并方法：选择性拼接、截断拼接和加权融合。每种方法都配有可直接运行的 PyTorch 代码，并分析其内存访问模式与计算开销。最后还会讨论工程实践中如何利用原地操作和索引技巧避免不必要的张量拷贝。

通道范数的计算与归一化处理

给定一个三维张量 x，形状为 (C, H, W)，它的通道范数可以写成一个长度为 C 的向量。以 L2 范数为例，计算方式是对每个通道 c 执行 torch.norm(x[c], p=2)，或者对整个张量在维度 (1,2) 上求范数：norms = x.norm(p=2, dim=(1,2))。这个操作的输出形状为 (C,)，每个元素对应一个通道的范数值。如果张量内存格式不是连续的，建议先调用 contiguous() 再计算，否则某些算子会触发隐式拷贝。另一种常见做法是使用全局平均池化得到通道均值作为重要性分数，但均值容易受到正负抵消影响，L2 范数则不存在这个问题，因此更适合作为通道筛选依据。

在合并两个张量之前，往往需要对范数做归一化，以便在不同分支之间公平比较。例如两个分支的数值尺度可能相差很大，一个分支的通道范数普遍是另一个分支的十倍。如果直接合并范数向量再排序，尺度大的分支会占据主导。归一化方式可以是最大最小值归一化，也可以是 softmax 归一化。实践中更稳妥的做法是分别对每个张量的通道范数做 L1 归一化，即除以所有通道范数之和，得到每个通道的相对重要性。这样两个张量的范数向量都处于同一量级，后续排序或加权才有意义。下面代码展示了计算通道范数并进行 L1 归一化的完整过程。

import torch

def channel_l2_norm(x):
    # x: (C, H, W)
    norms = x.norm(p=2, dim=(1, 2))
    return norms

def l1_normalize(norms):
    # 归一化到和为1，避免除零
    s = norms.sum()
    if s > 1e-8:
        return norms / s
    else:
        return torch.zeros_like(norms)

# 示例：两个形状不同的特征张量
x1 = torch.randn(32, 16, 16)
x2 = torch.randn(48, 16, 16)

n1 = channel_l2_norm(x1)
n2 = channel_l2_norm(x2)
n1_norm = l1_normalize(n1)
n2_norm = l1_normalize(n2)
print(n1_norm.shape, n2_norm.shape)

归一化之后，范数向量可以像概率分布一样使用，比如根据范数大小采样通道，或者作为加权系数。需要注意的是，L1 归一化会把所有范数之和压缩到 1，如果某个通道范数原本就接近零，归一化后依然接近零，不会放大噪声。但如果整个张量的所有通道范数都非常小，例如接近全零张量，那么归一化后的数值会不稳定，此时应加一个 epsilon 或者直接跳过该张量。

基于范数排名的选择性拼接

最直接的合并思路是：分别计算两个张量的通道范数，按范数从大到小排序，从两个张量中各选取前 k 个最重要的通道，然后把选出的通道拼接到一起。这样输出张量的通道数是 2k，而不是原始两个通道数之和。如果两个张量的原始通道数分别是 C1 和 C2，且 C1 和 C2 都大于 k，那么最终通道数固定为 2k，可控且高效。这种方法的优点是实现简单，缺点是只考虑各自内部排名，没有考虑两个张量之间相同位置通道的互补性，但通常已经能去掉大量低范数冗余通道。

实现排序选择时，可以用 torch.topk 直接得到范数最大的 k 个通道的索引，然后通过 index_select 或高级索引提取对应通道。这里推荐使用 torch.index_select，因为它在 GPU 上的内存访问更规整，比 Python 循环快得多。下面代码给出一个完整示例，其中 k 取两个张量通道数较小值的一半，也可以根据需求手动指定。

import torch

def merge_topk_by_norm(x1, x2, k):
    # x1, x2: (C, H, W)
    n1 = x1.norm(p=2, dim=(1, 2))
    n2 = x2.norm(p=2, dim=(1, 2))
    _, idx1 = torch.topk(n1, k)
    _, idx2 = torch.topk(n2, k)
    sel1 = torch.index_select(x1, 0, idx1)
    sel2 = torch.index_select(x2, 0, idx2)
    merged = torch.cat([sel1, sel2], dim=0)
    return merged

x1 = torch.randn(64, 8, 8)
x2 = torch.randn(96, 8, 8)
k = min(x1.size(0), x2.size(0)) // 2
result = merge_topk_by_norm(x1, x2, k)
print(result.shape)  # torch.Size([2*k, 8, 8])

这种选择性拼接的一个潜在问题是：如果两个张量本身通道数差异很大，比如一个 256 通道、一个 16 通道，而 k 取 8，那么小张量贡献 8 个通道、大张量也贡献 8 个通道，但大张量丢失了大量信息。这时可以按比例分配 k，比如从每个张量选取与其原始通道数成比例的通道数。以下代码展示了按比例选取的实现，它先根据范数总和确定两个张量的相对重要性，再计算各自的选取数量。

import torch

def merge_proportional_by_norm(x1, x2, total_k):
    n1 = x1.norm(p=2, dim=(1, 2)).sum()
    n2 = x2.norm(p=2, dim=(1, 2)).sum()
    total = n1 + n2
    if total < 1e-8:
        k1 = total_k // 2
        k2 = total_k - k1
    else:
        ratio1 = n1 / total
        k1 = int(round(total_k * ratio1))
        k2 = total_k - k1
    k1 = max(1, min(k1, x1.size(0)))
    k2 = max(1, min(k2, x2.size(0)))
    _, idx1 = torch.topk(x1.norm(p=2, dim=(1, 2)), k1)
    _, idx2 = torch.topk(x2.norm(p=2, dim=(1, 2)), k2)
    return torch.cat([torch.index_select(x1, 0, idx1),
                      torch.index_select(x2, 0, idx2)], dim=0)

按比例选取更能保留两个分支的原始信息分布，尤其适合多尺度特征融合场景，比如 FPN 中不同层级的特征图通道数不同，但都希望合并后保留足够细节。不过要注意，范数总和只是一个粗略的重要性指标，如果某个分支整体范数偏低但包含关键语义信息，按比例选取可能会被低估，此时可以结合任务先验调整比例。

加权融合与自适应合并策略

除了选择通道再拼接，另一种更灵活的方式是对两个张量的通道做加权求和。如果两个张量形状完全相同，即都是 (C, H, W)，那么可以直接按通道加权：merged = alpha * x1 + beta * x2，其中 alpha 和 beta 可以基于通道范数计算，也可以设计成可学习参数。基于范数的加权方式常见做法是：计算每个通道的范数，然后使用 softmax 生成两个分支的权重，这样每个通道都有自己的融合比例，而不是全局一个标量。具体来说，对于通道 c，权重 w1_c 和 w2_c 由 softmax([n1_c, n2_c]) 得到，然后 merged_c = w1_c * x1_c + w2_c * x2_c。这种合并保留了所有通道，但通过权重抑制低范数通道的贡献，相当于一种软性筛选。

实现按通道加权融合的代码并不复杂，直接利用广播机制即可。下面示例中两个张量形状相同，均为 (C, H, W)，计算通道范数后拼接成形状 (C, 2) 的矩阵，再沿最后一个维度做 softmax，得到两个分支的权重矩阵 shape (C, 2)，最后分别取权重列与原始张量相乘并相加。注意乘法时要把权重 reshape 成 (C, 1, 1) 以便广播到空间维度。

import torch
import torch.nn.functional as F

def weighted_merge_by_norm(x1, x2, temperature=1.0):
    # x1, x2: (C, H, W) 形状相同
    n1 = x1.norm(p=2, dim=(1, 2), keepdim=True)  # (C,1)
    n2 = x2.norm(p=2, dim=(1, 2), keepdim=True)
    # 拼接范数，形状 (C,2)
    norm_pair = torch.cat([n1, n2], dim=1) * temperature
    weights = F.softmax(norm_pair, dim=1)  # (C,2)
    w1 = weights[:, 0].view(-1, 1, 1)
    w2 = weights[:, 1].view(-1, 1, 1)
    merged = w1 * x1 + w2 * x2
    return merged

x1 = torch.randn(32, 16, 16)
x2 = torch.randn(32, 16, 16)
out = weighted_merge_by_norm(x1, x2)
print(out.shape)  # torch.Size([32, 16, 16])

如果希望合并后的张量通道数减少，可以在加权之后再做一次通道选择，或者引入可学习的线性变换将 2C 通道映射回 C 通道。但更常见的是将加权融合直接嵌入到网络结构中，让权重通过反向传播自动调整。此时可以初始化一个形状为 (C, 2) 的参数矩阵，并在前向传播中用 softmax 或 sigmoid 激活，这样权重既受初始范数引导，又能在训练中优化。以下代码展示了一个简单的可学习加权融合模块。

import torch
import torch.nn as nn

class LearnableNormMerge(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(channels, 2))
    def forward(self, x1, x2):
        # 初始时用范数作为先验设置权重
        with torch.no_grad():
            n1 = x1.norm(p=2, dim=(1, 2), keepdim=True)
            n2 = x2.norm(p=2, dim=(1, 2), keepdim=True)
            if self.weight.sum() == 0:
                init_w = torch.cat([n1, n2], dim=1)
                self.weight.copy_(init_w.squeeze(-1))
        w = torch.softmax(self.weight, dim=1)
        w1 = w[:, 0].view(-1, 1, 1)
        w2 = w[:, 1].view(-1, 1, 1)
        return w1 * x1 + w2 * x2

merge = LearnableNormMerge(32)
x1 = torch.randn(32, 16, 16)
x2 = torch.randn(32, 16, 16)
y = merge(x1, x2)
print(y.shape)  # torch.Size([32, 16, 16])

需要注意的是，上面的 if self.weight.sum() == 0 判断只适合权重全零初始化且尚未训练的场景，实际项目中建议在 __init__ 里直接根据输入张量初始化权重，避免在 forward 中做条件分支影响性能。加权融合的优点是输出通道数不变，不会丢失任何原始信息，缺点是计算量比选择性拼接略大，因为要对所有空间位置做两次乘法和一次加法。对于高分辨率特征图，这种开销不可忽略，因此需要根据实际资源选择方案。

内存优化与工程实践建议

在 GPU 上合并两个三维张量时，内存带宽往往是瓶颈。例如一个 256×128×128 的 float32 张量占用约 16MB，两个就是 32MB，如果每次合并都产生新的中间张量，很容易造成显存碎片和额外分配。优化手段之一是尽量使用原地操作，比如在加权融合中用 x1.mul_(w1).add_(x2 * w2) 替代 w1 * x1 + w2 * x2，这样可以复用 x1 的存储，减少一次临时分配。但要注意原地操作会破坏原始输入，如果后续还需要用到 x1 和 x2，就不能原地修改，需提前克隆一份。

另一个实用技巧是避免频繁调用 torch.norm 这类会分配输出的算子。如果只是需要范数排名，可以用 torch.sum(x * x, dim=(1,2)) 得到平方和，开根号只影响数值不改变排序，因此排序阶段可以不开根号，省去一次全通道的 sqrt 运算。对于大规模张量，还可以把两个张量先按通道维度 cat 成一个大张量，一次性计算所有通道范数，再分割索引，这样只调用一次底层 kernel，减少启动开销。下面代码演示了这种批量计算范数然后合并的做法。

import torch

def merge_batch_norm_topk(x1, x2, k1, k2):
    # 先拼接，一次性计算范数
    combined = torch.cat([x1, x2], dim=0)
    sq_norms = (combined * combined).sum(dim=(1, 2))  # 不开根号
    C1 = x1.size(0)
    n1, n2 = sq_norms[:C1], sq_norms[C1:]
    _, idx1 = torch.topk(n1, k1)
    _, idx2 = torch.topk(n2, k2)
    sel1 = combined[:C1][idx1]
    sel2 = combined[C1:][idx2]
    return torch.cat([sel1, sel2], dim=0)

这种批量计算方式在通道数较大时效果明显，因为减少了 kernel 调用次数。不过需要额外存储 combined 张量，如果 x1 和 x2 本身就很大，拼接后的临时张量会占用更多显存，此时可以改用分次计算。实践中还要注意张量的内存格式，若输入来自卷积层输出，通常是 channels_last 或 channels_first，不同格式下计算范数的效率差别很大。对于 channels_last 格式，在空间维度求和可能跨步访问，速度较慢，可先用 to(memory_format=torch.contiguous_format) 转换。最后，如果合并操作在网络训练中高频执行，建议将合并逻辑封装成 nn.Module 并注册到计算图中，避免每次手动计算范数造成 autograd 记录额外节点，影响反向传播性能。

总体而言，基于通道范数的三维张量合并是一个灵活且高效的特征融合手段。选择性拼接适合需要缩减通道数的场景，加权融合适合保持通道数并追求更平滑的信息整合，而工程优化则能进一步压缩时间与内存开销。根据实际任务选择合适的策略，往往能比盲目拼接获得更好的精度与速度平衡。

三维张量通道范数张量合并修改时间：2026-09-26 09:35:31