torch.unique与np.unique处理float64张量失效，求ε邻域去重方案

你遇到的确实是浮点数精度导致的问题——两个看似相同的数，实际存在极小的数值差异，被torch.unique或np.unique当成了不同元素。下面给出两种高效的解决方案，分别适配numpy数组和PyTorch张量：

一、Numpy 实现方案 #

先排序再筛选相邻元素差异，逻辑简单且高效：

import numpy as np

def unique_with_epsilon(arr, eps=1e-12):
    arr_sorted = np.sort(arr)
    # 计算相邻元素的差值
    diffs = np.diff(arr_sorted)
    # 第一个元素必保留，后续只留差值大于ε的元素
    keep_indices = np.concatenate([[True], diffs > eps])
    return arr_sorted[keep_indices]

排序后只保留和前一个元素差异超过ε的项，时间复杂度主要由排序决定（O(n log n)），适合大多数场景。

二、PyTorch 实现方案 #

和Numpy思路一致，适配张量场景还能利用GPU加速：

import torch

def unique_with_epsilon_torch(tensor, eps=1e-12):
    tensor_sorted, _ = torch.sort(tensor)
    diffs = torch.diff(tensor_sorted)
    # 构造掩码，第一个元素必保留
    keep_mask = torch.cat([torch.tensor([True], device=tensor.device), diffs > eps])
    return tensor_sorted[keep_mask]

如果你的张量在GPU上，这个方法能直接在显存中完成计算，不用来回拷贝数据。

三、更高效的量化替代方案 #

如果数据量极大，还可以先对数据做量化处理，再用普通的unique方法：

# Numpy版本
def unique_via_quantization(arr, eps=1e-12):
    # 把数值按ε粒度取整，相当于把ε邻域内的数映射到同一个值
    quantized = np.round(arr / eps)
    unique_quantized = np.unique(quantized)
    return unique_quantized * eps

# PyTorch版本
def unique_via_quantization_torch(tensor, eps=1e-12):
    quantized = torch.round(tensor / eps)
    unique_quantized = torch.unique(quantized)
    return unique_quantized * eps

这个方法的时间复杂度更低，但要注意量化会带来微小误差，适合对精度要求不是极端严格的场景。

注意点 #

ε的取值要根据数据精度调整：对于float64类型，选1e-12到1e-10比较合理，既能过滤精度误差，又不会把真正不同的数值误判为重复。

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