★ FEATURED ARTICLE 张量这玩意儿刚上手的时候大家注意力基本都放在形状变换、设备搬运、自动求导这几件事上等真正开始写训练循环才发现一半的代码都在跟统计量打交道loss 要取均值、梯度要算范数、特征要做归一化、评估指标要算准确率和方差。这些操作单个拎出来都不难难的是搞清楚哪个函数在哪个维度上动手、返回的是标量还是张量、会不会悄悄改变计算图。我见过不少项目模型结构写得漂漂亮亮结果因为一个dim参数写错训练半天指标纹丝不动排查了大半天才发现是统计维度搞反了。这篇内容围绕 PyTorch 里的张量统计计算展开把sum、mean、std、var、norm、argmax、cumsum这一批高频函数的使用逻辑、维度语义、计算图行为和实际场景里的取舍讲透。适合已经能跑通简单模型、但在数据处理和训练监控环节经常被统计操作绊住的读者也适合想系统梳理一遍基础 API 的人。下面这些内容大部分来自实际项目里的反复试错少数是踩坑之后回头翻文档才补上的原理我会尽量把为什么这么写和这么写会出什么问题都交代清楚。1. 统计函数的维度语义dim 到底在消掉谁1.1 从形状变化理解 dim 的作用张量统计函数里最容易出问题的参数就是dim。很多人对它的理解停留在沿着某个方向算这个说法不算错但太模糊真正写代码时需要的是更精确的模型指定 dim 之后这个维度会被压掉除非你显式保留它。拿一个形状为(B, C, H, W)的四维张量举例这是视觉任务里最常见的布局。执行x.mean(dim1)之后形状变成(B, H, W)通道维度消失了每个位置上的值是该位置所有通道的平均。如果写成x.mean(dim(2, 3))形状变成(B, C)空间维度被压平每个值是一张特征图上所有像素的平均。这个压掉的直觉比沿着方向好用得多因为你可以直接通过形状变化来验证自己写对没有。import torch x torch.randn(8, 3, 32, 32) print(x.mean(dim1).shape) # torch.Size([8, 32, 32]) print(x.mean(dim(2, 3)).shape) # torch.Size([8, 3]) print(x.mean(dim(1, 2, 3)).shape) # torch.Size([8])实际写代码时我习惯先在心里过一遍目标形状再反推 dim 该填什么。比如做全局平均池化目标是把(B, C, H, W)变成(B, C)那显然要压掉后两个维度dim(2, 3)就出来了。这种从目标形状反推的思路比死记硬背可靠得多尤其是在维度数量变化的时候。1.2 keepdim 参数的真实用途keepdimTrue这个参数经常被当成可有可无的选项实际上它在广播场景里是刚需。还是上面的例子x.mean(dim(2, 3), keepdimTrue)返回的形状是(B, C, 1, 1)而不是(B, C)。多出来的两个 1 看起来碍眼但它们让结果可以直接和原张量做广播运算。x torch.randn(8, 3, 32, 32) mean x.mean(dim(2, 3), keepdimTrue) # (8, 3, 1, 1) centered x - mean # 广播成立结果 (8, 3, 32, 32)如果不加keepdimx - mean会因为形状(8, 3, 32, 32)和(8, 3)对不上而报错或者在某些情况下触发意料之外的广播行为。做特征标准化、计算协方差、实现 LayerNorm 这类操作时keepdimTrue基本是标配。我的经验是只要统计结果还要参与后续的逐元素运算就默认加上 keepdim省得回头补 reshape。1.3 多维度统计的顺序问题dim传元组时多个维度是同时被压掉的不存在先后顺序的差异。这一点和连续调用两次单维度统计在数学上等价但计算路径不同。x.mean(dim(1, 2))是一次归约而x.mean(dim1).mean(dim1)是两次归约后者在数值上可能有微小差异而且中间结果的形状变化需要额外留意。提示涉及浮点累加的多步归约顺序不同可能带来精度差异。对精度敏感的场景比如大规模求和后再求均值优先用单次多维归约减少中间舍入。另外要注意dim传负数时是从后往前数的dim-1表示最后一维。这个写法在处理形状不固定的张量时特别方便比如序列模型里(B, T, D)的布局x.mean(dim-1)永远是对特征维求平均不用关心前面有几个维度。我在写通用工具函数时几乎都用负索引可读性和鲁棒性都更好。2. 均值、方差与标准差数值稳定性比公式更重要2.1 var 和 std 的默认行为差异torch.var和torch.std默认使用无偏估计也就是除以N-1而不是N。这个默认值来自统计学传统但在深度学习里经常不是你想要的行为。批归一化里的方差计算用的是有偏估计除以N如果你自己手写归一化逻辑却用了默认的var结果会和BatchNorm层对不上误差在小 batch 时尤其明显。x torch.randn(100) print(torch.var(x).item()) # 无偏除以 99 print(torch.var(x, unbiasedFalse).item()) # 有偏除以 100两者的差异在样本量大时可以忽略但当统计维度上的元素个数很少时比如小 batch 的通道统计差别能到百分之几。我的做法是手写归一化相关逻辑时显式写出unbiased参数不管选哪个至少让读代码的人一眼看到你的意图而不是靠默认值猜。2.2 手写标准化的正确姿势特征标准化是统计计算最典型的应用场景。给定一个(B, D)的特征矩阵按特征维做标准化的标准写法是mean x.mean(dim0, keepdimTrue) # (1, D) std x.std(dim0, unbiasedFalse, keepdimTrue) # (1, D) normalized (x - mean) / (std 1e-8)这里有几个细节值得说。dim0是因为要跨样本统计每个特征的分布这是按特征标准化的语义。unbiasedFalse是为了和主流归一化层的实现保持一致。分母加1e-8是防止除零这个 epsilon 的量级需要根据数据范围调整特征值本身很小的时候1e-8可能不够我遇到过特征量级在1e-6左右的数据集epsilon 得加到1e-5才稳。还有一个容易忽略的点如果x里存在常量特征某一列所有值相同std会是 0即使加了 epsilon归一化后的值也会被压到接近 0这个特征实际上就废了。生产环境里我会在标准化之前先检查一下有没有零方差列有的话要么剔除要么单独处理。2.3 用 Welford 思路理解流式统计当数据量大到无法一次性载入内存时均值和方差需要增量计算。PyTorch 本身没有直接暴露流式统计的 API但理解它的原理对排查数值问题很有帮助。直接累加平方和再减去均值平方的朴素做法在数值量级差异大时会出现严重的抵消误差这就是经典的灾难性抵消。Welford 算法通过维护运行均值M和运行平方差和S每来一个新样本就增量更新避免了直接计算大数相减。虽然日常用 PyTorch 的mean/var不需要自己实现但当你发现分批统计再合并的结果和一次性统计对不上时八成就是合并方式用了朴素公式。合并两组统计量的正确做法是按样本数加权并且方差项要加上组间均值差带来的修正项这个修正项经常被漏掉。3. 范数计算p 值选择背后的场景逻辑3.1 torch.norm 的 p 参数语义torch.norm的p参数决定了范数类型不同 p 值对应完全不同的几何意义和优化行为。p2 是欧几里得范数最常用梯度行为平滑p1 是绝对值之和对应稀疏性诱导pinf 是最大绝对值用于衡量最坏情况。x torch.tensor([3.0, -4.0]) print(torch.norm(x, p2).item()) # 5.0 print(torch.norm(x, p1).item()) # 7.0 print(torch.norm(x, pfloat(inf)).item()) # 4.0梯度裁剪是范数计算的高频场景。按范数裁剪的写法是先算全局范数再按比例缩放total_norm torch.norm(torch.stack([p.grad.norm(2) for p in model.parameters() if p.grad is not None]), p2) clip_coef max_norm / (total_norm 1e-6) if clip_coef 1: for p in model.parameters(): if p.grad is not None: p.grad.mul_(clip_coef)这里用torch.stack把所有参数的梯度范数拼成一个向量再求范数得到的是全局范数。注意p.grad.norm(2)里的 2 是 p 值不是维度别和dim参数搞混。另外clip_coef只在小于 1 时才生效这个判断不能省否则会把小梯度放大。3.2 范数与距离的区分新手常把范数和距离混为一谈。torch.norm(a - b)确实等于 a 和 b 之间的欧氏距离但torch.norm(a)单独拿出来是 a 到原点的距离语义完全不同。做相似度计算时如果直接对两个特征向量分别求范数再比较逻辑就错了正确做法是先做差再求范数或者用余弦相似度。余弦相似度的实现也依赖范数def cosine_sim(a, b, dim-1, eps1e-8): a_norm a / (a.norm(dimdim, keepdimTrue) eps) b_norm b / (b.norm(dimdim, keepdimTrue) eps) return (a_norm * b_norm).sum(dimdim)这里的eps同样是为了防止零向量导致除零。零向量在嵌入模型里并不罕见padding 位置、未激活的专家输出都可能是全零不加保护会直接产生 NaN然后污染整个 batch 的梯度。3.3 矩阵范数与谱范数的使用边界torch.norm对矩阵默认计算的是 Frobenius 范数也就是所有元素平方和开根号等价于把矩阵展平后求 L2 范数。如果要做谱范数最大奇异值需要显式指定p2配合dim(0, 1)或者直接用torch.linalg.norm的ord参数。A torch.randn(4, 5) print(torch.norm(A).item()) # Frobenius print(torch.linalg.norm(A, ord2).item()) # 谱范数 print(torch.linalg.norm(A, ordfro).item()) # 显式 Frobenius谱范数在对抗训练、 Lipschitz 约束、权重谱归一化里是核心工具但计算成本比 Frobenius 高不少因为它涉及奇异值分解或幂迭代。日常监控权重健康度用 Frobenius 就够了只有在做理论相关的约束时才需要上谱范数。我见过有人在训练循环里每步都算谱范数做监控结果训练速度掉了一大截其实完全没必要。4. 极值与索引类统计argmax 的坑比想象中多4.1 argmax 的返回语义与 tie 处理torch.argmax返回的是最大值的索引不是最大值本身。这个区别在分类任务里很关键模型输出 logits 之后argmax(dim1)得到的是预测类别而max(dim1)返回的是(values, indices)元组两个都要用的时候用后者更高效避免重复扫描。logits torch.randn(16, 10) pred logits.argmax(dim1) # (16,) values, indices logits.max(dim1) # 两个 (16,)遇到并列最大值时argmax返回的是第一个出现的索引这个行为在文档里有说明但经常被忽略。如果数据里有大量重复值比如量化后的输出、padding 掩码返回的索引可能不符合业务预期。需要稳定处理 tie 的场景要么在比较前加微小扰动要么自己写逻辑显式处理。4.2 topk 与排序的取舍需要前 k 个结果时用torch.topk它返回排序后的值和索引比先sort再切片高效。topk的largest参数控制取最大还是最小sorted参数控制是否保证返回结果有序。如果只关心集合不关心顺序把sorted设为False能省一点开销。scores torch.randn(1000) top5_vals, top5_idx torch.topk(scores, k5, largestTrue, sortedTrue)一个实际经验在检索或推荐场景里做 top-k 召回时如果 k 接近候选集大小topk的优势就不明显了直接全排序可能更简单。另外topk在 GPU 上的实现对大 k 值的效率不如小 k 值k 超过几千时值得测一下实际耗时再决定。4.3 cumsum 与累积统计的边界torch.cumsum做累积求和在计算累积分布、前缀和、序列位置编码时有用。它的输出形状和输入一致不会压掉维度这点和归约类函数不同。累积计算在长序列上要注意数值增长浮点累积到后面可能出现精度损失尤其是序列很长且元素量级不小的时候。x torch.ones(10000) c torch.cumsum(x, dim0) print(c[-1].item()) # 10000.0整数范围内没问题如果累积的是概率值或者小量级的浮点数长序列末尾的精度问题会更隐蔽。需要高精度累积时可以考虑用float64或者分段累积再合并。这个细节在实现累积分布函数、计算 AUC 这类指标时值得留意。5. 统计计算与计算图的交互detach 的时机5.1 统计量是否需要梯度这是实际项目里最容易出隐性 bug 的地方。张量统计函数默认会保留计算图也就是说loss.mean()是可导的梯度会回传到原始张量。但很多统计量只是用来做监控或日志不应该参与反向传播这时候必须detach。loss criterion(output, target) running_loss loss.detach().item() # 正确只取值不建图 # running_loss loss.item() # 也可以item 本身就脱离图loss.item()和loss.detach().item()在效果上都能拿到 Python 标量区别在于item()要求张量只有一个元素。如果统计的是多元素张量比如每个样本的 loss就得先mean()再item()或者detach().cpu().numpy()拿数组。注意在训练循环里对需要梯度的张量反复做统计而不 detach计算图会越积越大显存缓慢上涨最后 OOM。这个现象很隐蔽因为每一步看起来都正常只是显存曲线在爬。5.2 用统计量做归一化时的梯度路径有些统计量是要梯度的比如 LayerNorm 里的均值和方差它们参与前向计算梯度必须能回传。这种情况下不能 detach而且要注意统计维度和归一化维度的一致性。手写 LayerNorm 的简化版def manual_layer_norm(x, eps1e-5): mean x.mean(dim-1, keepdimTrue) var x.var(dim-1, unbiasedFalse, keepdimTrue) return (x - mean) / torch.sqrt(var eps)这里的mean和var都保留了计算图梯度会通过它们回传这正是 LayerNorm 梯度公式里包含均值项和方差项的原因。如果误加了detach梯度路径就断了训练会以非常隐蔽的方式退化——loss 能降但效果不对。5.3 分布式场景下的统计量同步多卡训练时单卡上的统计量只是局部统计全局统计需要跨卡归约。torch.distributed提供了all_reduce来做这件事但要注意归约的是和还是均值以及样本数是否均衡。local_sum x.sum() local_count torch.tensor(x.numel(), devicex.device, dtypetorch.float32) torch.distributed.all_reduce(local_sum) torch.distributed.all_reduce(local_count) global_mean local_sum / local_count直接对局部均值做all_reduce再除以卡数在各卡样本数不等时会算错。正确做法是归约分子和分母再相除这个细节在实现跨卡评估指标时特别重要我见过因为直接平均局部指标导致评估结果偏高的案例。6. 高频统计操作的性能与精度权衡6.1 归约操作的设备差异同样的mean操作在 CPU 和 GPU 上的性能特征不同。GPU 上的归约是并行实现的维度大小和内存布局会影响效率。对(B, C, H, W)做dim1的归约如果内存是通道优先布局访问模式就不如对最后一维归约友好。实际项目里如果统计操作成为瓶颈可以考虑调整张量布局把要归约的维度放到最后。x torch.randn(64, 256, 32, 32, devicecuda) # 对最后一维归约通常更快 y x.mean(dim-1)这个差异不是绝对的取决于具体硬件和 PyTorch 版本值得用torch.cuda.Event实测一下。我的习惯是在性能敏感的路径上对候选写法各跑几百次取中位数别凭直觉下结论。6.2 混合精度下的统计精度用autocast做混合精度训练时统计操作的行为需要留意。mean在float16下累加容易溢出或损失精度PyTorch 的 autocast 会把部分归约操作提升到float32但不是所有都提升。对精度敏感的统计比如 loss 累加、大张量求和显式转float32更稳妥。with torch.autocast(device_typecuda, dtypetorch.float16): out model(x) loss criterion(out, target) # 统计时显式提升精度 loss_f loss.float() batch_mean loss_f.mean().item()float16能表示的最大值约 65504大张量求和很容易超。float32的累加也有精度上限但日常规模够用。真正大规模求和比如百万级元素时可以考虑分块累加或者用float64代价是速度。6.3 统计函数的组合优化多个统计量需要同时计算时尽量复用中间结果。比如同时要均值和方差torch.var内部已经算了均值单独再调mean是重复计算。虽然 PyTorch 不会自动做这种公共子表达式消除但你可以手动组织# 不推荐两次扫描 m x.mean(dim0) v x.var(dim0, unbiasedFalse) # 推荐一次扫描拿两个统计量用 torch.var_mean v, m torch.var_mean(x, dim0, unbiasedFalse)torch.var_mean和torch.std_mean这类组合函数就是为这种场景准备的一次遍历同时返回两个统计量省一半的访存。在大张量上这个优化很实在尤其是统计维度很大的时候。7. 排查统计相关 bug 的实战思路统计计算出的问题往往不报错只是结果不对排查起来比形状错误麻烦。我总结了一套从外到内的检查顺序基本能覆盖大部分情况。第一步永远是打印形状。统计函数的输出形状是否符合预期是判断dim和keepdim有没有写错的最快方式。形状对了但值不对再往下查。第二步用小数据手算验证。构造一个 2x3 的小张量手动算出期望的均值、方差和函数结果对比。这一步能排除掉维度语义理解错误和unbiased参数用错这两类问题。第三步检查计算图状态。如果统计量参与了反向传播但结果异常用tensor.requires_grad和tensor.grad_fn确认图是否按预期构建。误 detach 和漏 detach 都会在这里暴露。第四步检查数值范围。NaN 和 Inf 是统计计算的常客来源通常是除零、溢出、log 负数。用torch.isnan和torch.isinf定位到具体位置再回溯是哪一步产生的。def check_tensor(name, t): if torch.isnan(t).any(): print(f{name}: 含 NaN位置 {torch.isnan(t).nonzero()[:5]}) if torch.isinf(t).any(): print(f{name}: 含 Inf位置 {torch.isinf(t).nonzero()[:5]}) print(f{name}: shape{tuple(t.shape)}, dtype{t.dtype}, fmin{t.min().item():.4f}, max{t.max().item():.4f})这个检查函数我在调试阶段会挂在关键节点上比事后翻日志高效得多。等逻辑稳定了再摘掉避免影响性能。还有一个经验统计维度上的元素个数为 1 时方差会是 0 或 NaN。unbiasedTrue且样本数为 1 时除以N-10直接产生 NaN。小 batch 或者某些维度被压到只剩一个元素时很容易触发加 epsilon 或者改用unbiasedFalse能规避。这个坑在 batch size 设为 1 做调试时几乎必踩提前知道能省不少时间。8. 几个容易被忽略的统计函数除了上面这些主力函数还有几个使用频率不高但特定场景下很好用的统计操作。torch.median返回中位数对异常值鲁棒在特征工程里做稳健统计时有用。它同样支持dim参数返回(values, indices)。中位数的计算成本比均值高因为需要部分排序大数据量下要权衡。torch.quantile计算分位数做数据分布分析、异常检测阈值确定时常用。它要求输入是浮点类型且分位数参数q的范围是[0, 1]。分位数在 GPU 上的实现效率一般大规模数据建议先采样再算。torch.mode返回众数离散数据统计时用得上。注意它返回的是(values, indices)且当有多个众数时只返回其中一个具体返回哪个没有明确保证依赖实现细节业务上如果需要全部众数得自己处理。torch.logsumexp是数值稳定的 log-sum-exp在计算 log 概率、softmax 的数值稳定版本、变分推断里的对数配分函数时是核心工具。它避免了先 exp 再 sum 再 log 的溢出问题内部做了最大值平移。手写 softmax 时用logsumexp能显著提升数值稳定性def stable_softmax(x, dim-1): return torch.exp(x - torch.logsumexp(x, dimdim, keepdimTrue))这些函数单独看都不复杂但知道它们在什么场景下比常规函数更合适是区分会用和用得好的地方。我的建议是先把mean、sum、var、norm、argmax这五个用熟遇到具体需求再查对应的专用函数不用一次性全记住。
阅读完成 · 觉得有帮助？
