PyTorch torch.unique() 基础与实战应用指南
作者：Geoking.
torch.unique() 是PyTorch中的一个去重函数,用于返回张量中所有的唯一元素（unique elements）,本文将带你深入了解 torch.unique() 的用法、参数、返回值以及实际应用场景,感兴趣的朋友跟随小编一起看看吧
在深度学习的数据处理中经常需要统计或筛选 张量（Tensor） 中的唯一值，比如去重、统计类别数量、计算唯一标签数等。
PyTorch 提供了一个非常方便的函数 —— torch.unique()，可以轻松完成这些操作。
本文将带你深入了解 torch.unique() 的用法、参数、返回值以及实际应用场景。
一、什么是torch.unique()？
torch.unique() 是 PyTorch 中的一个去重函数，用于返回张量中所有的唯一元素（unique elements）。
它类似于 Python 的 set() 或 NumPy 的 np.unique()，但专为 GPU 加速的张量操作 设计。
二、函数语法
torch.unique(input, sorted=True, return_inverse=False, return_counts=False, dim=None)
三、参数说明
| 参数 | 类型 | 说明 |
|---|---|---|
input | Tensor | 输入张量 |
sorted | bool | 是否对结果排序（默认 True） |
return_inverse | bool | 是否返回原张量中每个值在唯一值列表中的索引 |
return_counts | bool | 是否返回每个唯一值的出现次数 |
dim | int 或 None | 按指定维度去重，默认对整个张量去重 |
四、基本用法
🎯 示例 1：基础去重
import torch x = torch.tensor([1, 2, 2, 3, 3, 3]) unique_x = torch.unique(x) print(unique_x)
输出：
tensor([1, 2, 3])
✅ 结果去除了重复值，并自动排序。
🎯 示例 2：不排序
x = torch.tensor([3, 2, 1, 3, 2]) unique_x = torch.unique(x, sorted=False) print(unique_x)
输出：
tensor([3, 2, 1])
当 sorted=False 时，结果的顺序与首次出现的顺序一致。
五、返回索引与计数
🎯 示例 3：return_inverse
return_inverse=True 会返回一个索引张量，表示原张量中每个元素在唯一值（即新张量）中的位置。
x = torch.tensor([2, 1, 2, 3]) u, inv = torch.unique(x, return_inverse=True) print(u) print(inv)
输出：
tensor([1, 2, 3]) tensor([1, 0, 1, 2])
解释：
- 唯一值为
[1, 2, 3] - 原数组
[2, 1, 2, 3]中：- 第一个元素 2 → 索引 1
- 第二个元素 1 → 索引 0
- 第三个元素 2 → 索引 1
- 第四个元素 3 → 索引 2
🎯 示例 4：return_counts
return_counts=True 会返回每个唯一值出现的次数。
x = torch.tensor([1, 2, 2, 3, 3, 3]) u, counts = torch.unique(x, return_counts=True) print(u) print(counts)
输出：
tensor([1, 2, 3]) tensor([1, 2, 3])
表示：
- 值 1 出现 1 次
- 值 2 出现 2 次
- 值 3 出现 3 次
🎯 示例 5：同时返回多个结果
你可以同时返回 unique 值、inverse 索引和计数：
x = torch.tensor([1, 2, 2, 3, 3, 3]) u, inv, counts = torch.unique(x, return_inverse=True, return_counts=True) print(u) print(inv) print(counts)
输出：
tensor([1, 2, 3]) tensor([0, 1, 1, 2, 2, 2]) tensor([1, 2, 3])
六、按维度去重（dim 参数）
默认情况下，torch.unique() 会将张量展开成一维后去重。
但如果你希望在特定维度上去重（如按行或按列），可以使用 dim 参数。
🎯 示例 6：按行去重
x = torch.tensor([[1, 2],
[1, 2],
[3, 4]])
unique_rows = torch.unique(x, dim=0)
print(unique_rows)输出：
tensor([[1, 2],
[3, 4]])
表示第 1、2 行重复，只保留一个。
🎯 示例 7：按列去重
x = torch.tensor([[1, 1, 3],
[2, 2, 4]])
unique_cols = torch.unique(x, dim=1)
print(unique_cols)输出：
tensor([[1, 3],
[2, 4]])
七、torch.unique()与 NumPy 对比
| 功能 | PyTorch (torch.unique) | NumPy (np.unique) |
|---|---|---|
| 默认排序 | ✅ 是 | ✅ 是 |
| 支持 GPU | ✅ 是 | ❌ 否 |
| 返回 inverse 索引 | ✅ 是 | ✅ 是 |
| 返回 counts | ✅ 是 | ✅ 是 |
| 按维度去重 | ✅ 是（dim） | ❌ 不直接支持 |
| 性能 | 高（GPU 支持） | 仅 CPU |
八、实际应用场景
1. 分类问题中统计类别数量
labels = torch.tensor([0, 1, 0, 2, 2, 1, 3])
classes = torch.unique(labels)
print(f"共有 {len(classes)} 个类别: {classes.tolist()}")
输出：
共有 4 个类别: [0, 1, 2, 3]
2. 计算样本分布（类别频率）
labels = torch.tensor([0, 1, 0, 2, 2, 1, 3])
u, counts = torch.unique(labels, return_counts=True)
for c, cnt in zip(u.tolist(), counts.tolist()):
print(f"类别 {c}: {cnt} 个样本")
输出：
类别 0: 2 个样本 类别 1: 2 个样本 类别 2: 2 个样本 类别 3: 1 个样本
3. 在图像分割中统计像素类别
例如在语义分割任务中，计算 mask 图像中有多少个不同的像素类别：
mask = torch.randint(0, 5, (256, 256)) # 随机生成类别标签
num_classes = len(torch.unique(mask))
print(f"图像中共有 {num_classes} 个类别")
⚠️ 九、注意事项
torch.unique()** 默认会对结果排序**，如果在意性能，可以设置sorted=False。- 对高维张量使用
dim去重时，必须保证该维度的所有元素形状一致。 - 对大张量使用
return_counts或return_inverse时可能会消耗更多显存。
📚 参考资料
到此这篇关于PyTorch torch.unique() 基础与实战应用指南的文章就介绍到这了,更多相关PyTorch torch.unique() 使用内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！