干这行久了你会发现很多看似高级的模型操作落到代码层面无非是把数据重新摆一摆、数一数、分组求和。PyTorch 里最容易被低估的三个统计函数——bucketize、bincount、histogram正是干这三件事的利器。它们单个拿出来都很简单但组合起来能覆盖特征分箱、标签统计、分布直方图、蒙特卡洛区间概率等一大半日常开发场景。这篇内容适合正在做数据处理、特征工程、模型评估的人也适合刚开始用 PyTorch、想摆脱“手写 for 循环统计”的小白。我会直接把它们揉进案例里讲你不仅能看懂每个参数怎么用还能知道为什么在某个场景该选它而不是另两个。下面涉及代码均在 PyTorch 2.x 的 CPU/GPU 环境验证过低版本个别 API 有差异我会在对应位置标出来。1. 为什么是这三兄弟先弄清每个函数到底是干嘛的1.1 bucketize瞬间完成连续值到区间编号的映射torch.bucketize(input, boundaries)的作用一句话就能说清给你一个有序的分界点列表 boundaries它返回每个输入元素落在哪个区间。注意返回的是索引范围从 0 到 len(boundaries)是一个整数张量。做年龄分箱、价格分段、特征离散化的时候这函数真能省掉一大段 if-else。我最早是这么干的bucket [] for x in input: for i, b in enumerate(boundaries): if x b: bucket.append(i) break后来有了 bucketize直接一行bucket_idx torch.bucketize(input, boundaries)顺带提一个冷知识torch.bucketize其实是torch.searchsorted的别名你搜 searchsorted 也能搜到同一套实现。它的核心是二分查找时间复杂度 O(n log m)百万级数据也就是毫秒级的事。因为返回的是整数索引后面可以直接接torch.bincount、torch.gather、torch.scatter_add非常顺手。正是这种“先映射成桶号再继续向后传递”的设计让它成为组合统计的第一步。1.2 bincount一维整数标签的频次计数器torch.bincount(input)是标准的频次统计函数输入一个非负整数张量输出一个长度等于“最大值加 1”的张量每个位置存数量。举个例子输入 [3, 3, 1, 0]输出就是 [1, 1, 0, 2]因为 0 出现一次、1 出现一次、2 没出现、3 出现两次。这个函数最妙的地方在于完全张量化能直接在 GPU 上跑。等于把 Python 端的 Counter 或 for 循环整个搬进了张量计算的世界里。我做过一次对比统计 100 万个整数标签的频次for 循环方案要几百毫秒bincount 在 GPU 上不到 1 毫秒差距是数量级的。bincount 还有两个非常实用的参数weights 和 minlength。weights 允许你对每个样本加权后再累加比如 KNN 投票时希望“近的样本权重大”就不用先数次数再挨个乘权重了minlength 则保证输出至少有多长比如统计 5 个类别的分布即使某个类一个样本都没有输出也是固定长度 5不会因为最大标签是 2 就输出长度 3。这个细节在拼模型特征时特别重要。1.3 histogram带边界信息的完整分布画像torch.histogram(input, bins, range)是 PyTorch 官方提供的直方图实现返回两个值直方图计数 hist 和每个区间的边界 bin_edges。它和 bincount 最大的区别是直接接收连续浮点数自动完成“分箱 计数”两件事。bins 可以传一个整数表示等宽分箱也可以传一个一维张量表示自定义边界。比如你想把数据分成 [0, 10)、[10, 50)、[50, 100) 三个不均匀的区间直接传torch.tensor([0., 10., 50., 100.])就行不用自己算边界索引。这特性在业务里很实用因为真实场景的分箱极少是等宽的。它还支持 weight 和 density 两个参数。weight 和 bincount 的套路一致densityTrue时输出的是概率密度也就是每个桶的计数除以总样本数再除以桶宽适合直接和概率密度曲线做对比。因为返回了边界张量后续画图、区间查表、分箱统计都方便。直方图这个统计概念被完整搬进张量计算后最大的收益就是你不用再把数据拷到 CPU 上用 numpy 和 matplotlib 来回倒腾了。2. 三个真实案例从特征分箱到梯度直方图2.1 案例一用户年龄分箱统计三函数串成一条流水线假设你现在有 10 万条用户年龄数据想按[0,18)、[18,25)、[25,35)、[35,45)、[45,60)、[60,)分成 6 组统计每组人数和占比。这是特征工程里最常见的操作直接用组合拳import torch # 模拟年龄数据均值35、标准差10的正态分布再截断到 [0, 80] ages (torch.randn(100000) * 10 35).clamp(0, 80) # 分界点 boundaries torch.tensor([18., 25., 35., 45., 60.]) # 第一步把每个年龄映射到桶编号0~6 bucket_idx torch.bucketize(ages, boundaries) # 第二步统计每个桶的数量minlength 强制输出7个位置 counts torch.bincount(bucket_idx, minlengthlen(boundaries) 1) print(counts)这里有个关键点要说明bucketize返回的索引范围是0到len(boundaries)所以当boundaries有 5 个元素时索引最大是 5合起来共 6 个桶。索引 5 的含义是“年龄大于等于 60”。如果你漏了minlengthlen(boundaries) 1bincount 的输出长度会根据实际出现过的最大索引调整但万一某个边界的桶是空的长度就会不一致后面拼到特征向量里直接出 bug。我用这行代码检查一下与histogram的差异hist, edges torch.histogram(ages, binsboundaries) print(hist) print(edges)hist的结果和上面的counts几乎一一对应因为torch.histogram在边界张量传入时就是按同样的左闭右开逻辑在分。区别在于histogram直接帮你把边界也返回了而bucketize bincount是两步走好处是你能拿到“每个样本属于哪个桶”这件事本身。2.2 案例二图像梯度方向直方图HOG 雏形目标检测和特征描述子里有个经典统计场景计算图像梯度方向直方图。传统 OpenCV 写法要嵌套循环PyTorch 加起来只需要几行。我们先粗略算一遍梯度方向和幅值再做方向分箱import torch import torch.nn.functional as F # 模拟一张灰度图 gray torch.randn(256, 256) # sobel 算子 sobel_x torch.tensor([[-1., 0., 1.], [-2., 0., 2.], [-1., 0., 1.]]).view(1, 1, 3, 3) sobel_y sobel_x.transpose(2, 3) gx F.conv2d(gray.view(1, 1, 256, 256), sobel_x, padding1) gy F.conv2d(gray.view(1, 1, 256, 256), sobel_y, padding1) # 梯度方向角度弧度制转到 [0, 360) angle torch.atan2(gy, gx) 3.14159265 angle_deg (angle * 180 / 3.14159265) % 360 # [0, 360) # 方式一直接用 histogram 统计9个方向区间 hist, edges torch.histogram(angle_deg, bins9, range(0., 360.)) print(hist) # 方式二用 bucketize bincount 拿到每个像素的方向桶编号 bin_idx torch.bucketize(angle_deg, torch.linspace(0., 360., 10)[:-1]) each_pixel_bin torch.bincount(bin_idx.view(-1), minlength9)方式一适合只要一个全局统计结果方式二适合你后面还要用桶编号做像素级特征拼接的情况比如每个像素要额外记录“它落在哪个方向区间”bucketize 返回的索引可以直接在后续计算里复用。这里也提醒一个细节torch.linspace(0., 360., 10)[:-1]生成的是 9 个左边界点。也就是说角度正好等于 40 度时会被分到下一个桶。2.3 案例三蒙特卡洛模拟后验概率的快速计算贝叶斯统计里经常要算“采样出来的样本落在某个区间的概率”。比如你跑完 10 万次模拟得到一组后验样本要求落在[-1,0)、[0,1)、[1,2)、[2,∞)的概率。这时候 bucketize 加 bincount 简直就是为这个场景准备的samples torch.randn(100000) boundaries torch.tensor([-1., 0., 1., 2.]) idx torch.bucketize(samples, boundaries) # minlength5 因为边界有4个桶编号范围是0~4 counts torch.bincount(idx, minlength5) probs counts.float() / samples.numel() print(probs)再看 histogram 的做法hist, edges torch.histogram(samples, binsboundaries) probs_hist hist / samples.numel() print(probs_hist)两种结果完全一致。这里有个意外收获torch.histogram支持用户自定义边界张量不需要等宽分箱这在处理非对称区间时太香了。以前用 numpy 时非等宽区间要把边界单独维护一份现在直接在张量里算完后续接梯度反向传播也完全没障碍。3. 细节与性能这些参数和边界策略我踩过坑3.1 边界开闭规则right、range、bins 的语义这三个函数的边界规则是同一个套路默认左闭右开。也就是说输入值v落在区间[a[i-1], a[i])时返回索引i。等于边界值的元素会被放进右边的桶。但bucketize提供了一个rightTrue参数可以把规则变成左开右闭(a[i-1], a[i]]。这两个规则谁更重要我建议你记住一条先确认业务上“边界归谁”。比如阈值判断“低于 18 岁”和“18 岁及以下”是完全不同的两套逻辑。直接用right参数控制比事后给数据加减一个极小量要靠谱得多。torch.histogram的区间语义也类似最后一个桶两边都是闭区间这是为了覆盖最大值。这个细节容易忽略因为绝大多数数据不会恰好卡在最大值上。说个实际例子。我有一回统计点击率特征的分桶分布边界里有一个值是 0结果等于 0 的样本被全部归到了非零桶导致两个相邻桶的分布看起来严重不均衡。排查了半天最后就是边界开闭规则的问题。从那以后所有涉及边界值的统计任务我都会先写个 5 行自测用例验证规则再跑正式数据。3.2 类型、设备与确定性GPU 上的隐藏差异这三个函数在 GPU 上都能跑但有几个隐藏差异值得展开。第一个是bincount的浮点权重在 GPU 上累加顺序不确定。如果 weights 是浮点数不同线程累加的顺序可能不同导致结果和 CPU 版本有极小的差异。绝大多数业务场景无所谓但你要是做精度复现实验比如训练结果要严格对齐就得留意。PyTorch 提供了torch.use_deterministic_algorithms(True)但开了之后某些算子会变慢甚至直接报错我的建议是只在复现实验里临时开并在确认结果一致后关掉。第二个是内存问题。bucketize默认返回 int64也就是 8 字节一个元素。如果是 1000 万样本一个索引张量就要 80MB 内存后面再接bincount会翻倍。此时可以传out_int32True直接省一半内存。代价是索引最大值不能超过 21 亿左右正常场景都没问题但如果你在跑上亿规模的数据还是老老实实用 int64。第三个是数据类型限制。torch.histogram要求输入必须是浮点张量输入整数张量会直接报错。这不算 bug直方图算法的边界计算需要浮点除法所以遇到整数输入先.float()再传。bincount则相反要求输入必须是整数传浮点会报错。这两个限制我一开始经常搞混现在总结成一句话想看分布用 histogram浮点进想数个数用 bincount整数进。3.3 新旧版本差异histc、histogram 和其他 API 的迁移如果你维护的老项目里见过torch.histc那就是torch.histogram的前身。torch.histc接口很简单torch.histc(input, bins100, min0, max0)。它的问题是只支持整数 bins 和 min/max 参数没有 density也没有自定义边界官方文档也已经建议优先使用torch.histogram。老代码迁移时注意min和max都传 0 表示“自动根据数据范围确定”迁移到 histogram 时直接省略range参数即可。还有torch.searchsorted的问题。前面说过 bucketize 是它的别名二者参数完全一致不需要纠结用哪个。不过如果代码里同时出现torch.bucketize和torch.searchsorted建议统一成一个避免团队里有人搞混。另外一个边界情况自定义边界张量传给torch.histogram时不同 PyTorch 版本对 GPU 的支持不完全一致。我实测下来torch 2.1 之后在 CUDA 上直接传边界张量已经没问题了但更早版本可能要求你先在 CPU 上算完再拷回 GPU。稳妥做法是写个 50 行小脚本在自己的环境里验证一次别直接信文档。4. 何时用哪个三函数的选型对比4.1 输入维度与输出形态对比这里直接放一张对照表是我平时最常用的选型依据函数输入要求返回内容典型场景bucketize连续值或整数任意形状与输入同形状的整数索引分箱、分段、特征离散化、区间映射bincount非负整数一维一维频次张量标签计数、加权投票、类别频次统计histogram浮点数任意形状频次张量 边界张量直方图、分布画像、非等宽分箱统计一个很容易踩的坑是维度。bincount只接受一维输入二维以上要先.view(-1)或.reshape(-1)。bucketize和histogram都支持任意形状输入前者返回同形状索引后者返回一维直方图结果。如果你拿二维特征矩阵直接调bincount报错信息会提示你“expected 1D tensor”不怕压平就行。4.2 和 numpy 的等价映射迁移时的对照经验从 numpy 生态迁到 PyTorch 的人通常会找等价函数。这里我给出对应关系np.searchsorted对应torch.bucketizenp.bincount对应torch.bincountnp.histogram对应torch.histogram。接口不能说一模一样但也足够接近迁移成本很低。需要特别小心的是searchsorted的边界语义。numpy 的side参数和 PyTorch 的right参数并不是完全一对一的直观关系。np.searchsorted(a, v, sideleft)返回的是“最左边的插入位置”而torch.bucketize(v, a, rightFalse)在某些边界情况下和它等价但我不建议盲目照搬。最稳妥的方式是构造一组包含边界值的测试数据比如torch.tensor([0., 1., 1., 2.])配boundariestorch.tensor([1.])分别跑一遍看输出是否一致。这种一次性验证能帮你省掉后续排查的半天时间。另外numpy 的bincount和 PyTorch 的bincount行为几乎完全一致都是输出长度等于最大标签加一。但 PyTorch 版本对空输入的处理要小心空张量输入且minlength0时输出也是空张量而 numpy 会报错或者返回空数组细节有差异。构造训练数据时我都会给minlength一个明确值确保数据为空时统计结果不会让下游逻辑雪崩。5. 常见问题与排查实录5.1 问题速查表问题现象原因解决方案bincount报错 “input tensor should be nonnegative”输入包含负数先做偏移labels labels - labels.min()bincount输出长度不稳定最大标签不同或没有指定 minlength统一传入minlength类别数histogram报错 “input tensor should be floating point”输入是整数张量提前.float()bucketize结果比预期多一个桶没理解“返回 0 到 len(boundaries)”的语义把 len(boundaries) 视为“超出上界”的桶边界附近样本在相邻桶间跳动频次忽高忽低数据存在 NaN 或轻微浮点误差先torch.isnan清洗再对边界做精确判定GPU 与 CPU 统计结果有微小差异浮点累加顺序不确定复现实验开启确定性算法常规任务忽略自定义边界张量报错“must be monotonically increasing”边界未排序或有重复torch.sort排序后torch.unique去重5.2 一次真实的排查过程消失的 3 条样本有一次我在做特征分布校验发现 bucketize 后所有桶的频次加起来比原始样本总数少了 3 条。第一反应是索引越界被截断了检查了 minlength没问题检查了 bincount 的输入也没问题。最后用torch.isnan(ages).sum()一查发现数据里有 3 个 NaN。NaN 在比较运算里的行为比较特殊任何和 NaN 比较的操作都返回 False。落到bucketize上NaN 会被排到最后一个桶但它本身不参与正常的数值比较导致统计结果在一些版本里会显得“少了一部分样本”或者说被归到了边界之外但你很难察觉。这事之后我给自己定了个规矩凡是交给这三个函数的数据先做一步清洗把 NaN 替换成哨兵值或者直接剔除。这个习惯现在一直保留着。另一个类似的问题是极端值。如果输入里有一个极大的值比如 1e9bucketize 会把它归到最后一个桶但 histogram 在range(min, max)的默认逻辑下会把最大值单拎出来当成最后边界。如果不提前看数据范围你会看到直方图尾部出现一个异常的尖峰。解决方案很简单分析前分别打印min、max、isnan数量做到心里有数。我个人在实际操作中的体会是这三个函数的价值不在于单个功能有多复杂而在于它们能无缝串成一条张量流水线。先想清楚边界怎么定再想清楚要返回索引还是返回直方图最后顺手把 NaN 和维度问题处理掉后面就基本不会出大麻烦。最后再分享一个小技巧当你需要反复统计不同边界分组时可以把boundaries和minlength先封装成一个配置对象保证整个项目里所有分箱口径完全一致。这样既避免了同事之间各写各的造成口径混乱也方便后期调整分箱边界时一处修改全局生效。这三个函数虽然不起眼但把它们的参数约定固定下来你在特征工程上的效率能提升得非常明显。