本文介绍了一种在 PyTorch 张量中高效查找每个唯一行首次出现索引的方法。通过利用 torch.unique 函数获取唯一行及其逆向索引，并结合二维张量和 torch.argmin 函数，避免了显式循环，从而提升了代码效率。文章提供了详细的代码示例和性能注意事项，帮助读者根据实际应用场景选择合适的解决方案。
在 PyTorch 中处理张量数据时，经常需要查找唯一行的索引。一种常见的方法是使用循环遍历每个唯一行，并在逆向索引中找到其首次出现的索引。然而，这种方法效率较低，尤其是在处理大型张量时。本文介绍一种更高效的方法，利用 PyTorch 的张量操作避免显式循环，从而提高代码性能。
使用 torch.unique 获取唯一行和逆向索引
首先，使用 torch.unique 函数获取张量中的唯一行、逆向索引和计数。torch.unique 函数的 return_inverse=True 参数会返回一个逆向索引张量，该张量指示原始张量中的每一行对应于唯一行张量中的哪个索引。
import torch import numpy as np # 示例张量 data = torch.rand(100, 5) data[np.random.choice(100, 50, replace=False)] = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) # 查找唯一行 u_data, inverse_indices, counts = torch.unique(data, dim=0, return_inverse=True, return_counts=True)
使用二维张量和 torch.argmin 查找首次出现索引
为了避免循环，我们可以创建一个二维张量 A，其维度为原始张量的行数乘以唯一行的数量。将 A 初始化为一个较大的值（例如 1000，确保大于原始张量的行数），表示“未定义的行索引”。然后，对于原始张量的每个行索引 i，将 A[i, inverse_indices[i]] 设置为 inverse_indices[i]。
A = 1000 * torch.ones((len(data), len(u_data)), dtype=torch.long) A[torch.arange(len(data)), inverse_indices] = inverse_indices
现在，考虑按列查看张量 A。第 j 列对应于第 j 个唯一行。该列的大部分值为 1000，但某些行将包含 j。该列的 argmin 就是映射到唯一行 j 的第一个原始行的索引。
unique_indices2 = torch.argmin(A, dim=0)
完整代码示例
import torch
import numpy as np
# 示例张量
data = torch.rand(100, 5)
data[np.random.choice(100, 50, replace=False)] = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
# 查找唯一行
u_data, inverse_indices, counts = torch.unique(data, dim=0, return_inverse=True, return_counts=True)
# 使用循环查找首次出现索引（作为参考）
unique_indices = torch.zeros(len(u_data), dtype=torch.long)
for idx in range(len(u_data)):
unique_indices[idx] = torch.where(inverse_indices == idx)[0][0]
# 使用二维张量和 argmin 查找首次出现索引
A = 1000 * torch.ones((len(data), len(u_data)), dtype=torch.long)
A[torch.arange(len(data)), inverse_indices] = inverse_indices
unique_indices2 = torch.argmin(A, dim=0)
# 验证结果
print(torch.allclose(unique_indices2,unique_indices))性能注意事项
虽然这种方法避免了循环和 torch.where 函数，但它使用了更多的内存。argmin 函数在硬件上的速度、实际问题的维度以及对内存的重视程度都会影响其效率。在实际应用中，需要根据具体情况权衡内存使用和计算速度，选择最合适的解决方案。如果数据量较小，循环方式可能更简单易懂；如果数据量较大，且对性能要求较高，则可以考虑使用本文介绍的基于张量操作的方法。
