如何实现PyTorch张量的批量直方图计算？

云上老周

2026-10-2

PyTorch 批量计算张量直方图实现方法

PyTorch 没有原生提供直接的batch_histogram接口，但可以通过现有算子快速实现符合你需求的功能，以下是两种常用的实现方案：

方案1：one_hot 向量化实现（推荐，GPU 性能更高）

该方案无循环，完全向量化，适合大部分批量计算场景：

import torch
import torch.nn.functional as F

def batch_histogram(x, bins=256, min=0, max=255):
    # 输入x要求batch维度在第0位，支持任意(B, ...)形状的输入
    batch_size = x.shape[0]
    # 拉平每个batch样本的元素为二维：(B, 单样本元素总数)
    x_flat = x.reshape(batch_size, -1)
    # 截断值到[min, max]区间，避免索引越界
    x_clamped = x_flat.clamp(min=min, max=max)
    # 计算每个值对应的bin索引
    bin_idx = ((x_clamped - min) / (max - min) * (bins - 1)).long()
    # one_hot编码后按样本维度求和得到直方图
    return F.one_hot(bin_idx, num_classes=bins).sum(dim=1)

调用方式和你给出的示例完全一致：

# 输入x形状为(64, 224, 224)
x = torch.randint(0, 256, (64, 224, 224)).cuda()
# 输出x形状为(64, 256)
x = batch_histogram(x, bins=256, min=0, max=255)

方案2：基于torch.histc的循环实现

适合batch尺寸小、单样本元素量极大的场景：

import torch

def batch_histogram(x, bins=256, min=0, max=255):
    batch_size = x.shape[0]
    hist = torch.zeros(batch_size, bins, device=x.device)
    for i in range(batch_size):
        hist[i] = torch.histc(x[i], bins=bins, min=min, max=max)
    return hist