PyTorch循环优化提速：非向量化方案及向量化思路问询

1. 无需向量化的循环提速方案 #

• 优化JIT编译逻辑：保留@torch.jit.script的同时，将Python列表probs替换为PyTorch张量（比如torch.zeros(len(t_samples), dtype=torch.int32, device=t_samples.device)），避免Python列表的动态修改开销；同时把torch.unbind的迭代方式换成索引遍历（for i in range(t_samples.shape[0])，通过t_samples[i]获取元素），JIT对索引遍历的优化稳定性更强。
• 减少跨层交互开销：避免在循环中触发Python-Torch的同步操作，比如原代码中if is_same is True会强制同步张量到CPU，可改为if is_same.all().item()（仅在必要时同步）；同时确保所有张量都在同一设备（CPU/CUDA）上，避免循环内的设备迁移。
• 利用Torch内置并行工具：CPU环境下通过torch.set_num_threads()设置匹配CPU核心数的线程数，让Torch内部操作自动多线程并行；GPU环境下优先使用CUDA张量，且尽量减少循环内的小张量创建，复用已有张量内存。
• 用torch.vmap实现自动批量映射：torch.vmap可以将单样本操作自动扩展为批量操作，无需手动改循环逻辑。比如把内层比较逻辑封装成单样本函数，通过vmap映射到整个样本集，性能接近手动向量化，代码改动极小。

2. 向量化优化的具体思路（以示例代码为例） #

针对基于容差统计样本出现次数的需求，推荐以下向量化实现路径：

广播+全局比较（最直接的方案） #

利用PyTorch的广播机制，一次性完成所有样本对的比较，再聚合结果：

def vectorized_seq_prob(t_samples: torch.Tensor):
    # 扩展维度实现样本对的广播比较
    t_i = t_samples.unsqueeze(1)  # 形状变为 [N, 1, D]
    t_k = t_samples.unsqueeze(0)  # 形状变为 [1, N, D]
    # 逐元素容差比较后，按样本维度聚合判断是否完全匹配
    match_matrix = torch.all(torch.isclose(t_i, t_k, rtol=1e-05, atol=1e-08), dim=-1)
    # 每行求和得到对应样本的匹配次数
    count = match_matrix.sum(dim=1)
    return count.tolist()

该方案完全消除Python循环，所有操作在Torch的C++/CUDA层执行，性能提升显著，逻辑也相对清晰。

量化哈希优化（针对低维样本） #

如果样本维度D较小，可通过量化将近似样本映射为相同哈希值，再统计次数：

1. 按容差atol对样本元素进行量化（比如torch.round(t_samples / atol) * atol）；
2. 将量化后的张量转换为一维哈希键（比如展平后用加权求和生成唯一标识，或直接用torch.unique的return_inverse参数）；
3. 用torch.bincount统计每个哈希键的出现次数，再映射回原样本。
   此方案的计算量远低于全局比较，适合样本量极大但维度低的场景。

分组去重优化 #

先提取近似唯一的样本组，再统计每组的出现次数，避免重复计算：

1. 自定义近似去重逻辑（通过torch.isclose结合广播找到唯一样本）；
2. 对每个唯一样本，统计原样本集中与其近似匹配的数量；
3. 将结果映射回原样本列表。
   该方案适合样本集中存在大量重复近似样本的场景，能大幅减少计算量。

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