PyTorch张量唯一行首次出现索引的高效查找方法

1. 问题背景与传统方法

在数据处理和机器学习任务中，我们经常需要处理包含重复数据的张量（tensor）。当需要识别张量中所有唯一行，并进一步获取这些唯一行在原始张量中首次出现的索引时，一个常见的挑战是效率问题。

PyTorch提供了torch.unique函数来方便地找出张量中的唯一行及其相关信息。例如，torch.unique(data, dim=0, return_inverse=True)会返回唯一行、以及一个inverse_indices张量，该张量将原始张量中的每个行映射到其对应的唯一行索引。

然而，要根据inverse_indices找出每个唯一行在原始张量中首次出现的索引，一个直观但效率低下的方法是使用Python循环：

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]

print("传统方法得到的首次出现索引：", unique_indices)

上述代码中，for循环遍历每个唯一行的索引idx，然后使用torch.where查找inverse_indices中所有等于idx的位置，并取第一个位置作为首次出现的索引。这种逐个查找的循环方式，尤其是在处理大型张量时，会导致显著的性能开销，因为它涉及多次Python循环迭代和张量条件查找操作。

2. 优化方法：基于二维张量和argmin的向量化方案

为了避免上述低效的循环，我们可以采用一种更符合PyTorch风格的向量化方法。其核心思想是构建一个辅助的二维张量，巧妙地利用其结构，并通过torch.argmin操作来高效地找出首次出现的索引。

核心思路：

1. 创建辅助张量A： 构建一个维度为 (原始行数, 唯一行数) 的二维张量A。将其所有元素初始化为一个足够大的占位符值（例如，远大于原始行数的整数）。
2. 填充张量A： 利用高级索引，将原始张量中的行索引映射到其对应的唯一行索引。具体来说，对于原始张量中的每一行i，如果它属于唯一行组j（即inverse_indices[i] == j），则在张量A的 (i, j) 位置填充值 i。
3. 使用argmin查找： 对张量A沿唯一行维度（dim=0，即列方向）执行torch.argmin操作。对于每一列j，argmin将返回该列中最小值所在的行索引。由于我们填充的值是原始行索引i，并且占位符值远大于任何有效的i，因此argmin将准确地找到属于唯一行组j的最小原始行索引，这正是我们所需的首次出现索引。

示例代码：

纳米搜索

纳米搜索：360推出的新一代AI搜索引擎

30

查看详情

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)

# 优化方法：基于二维张量和argmin
num_original_rows = len(data)
num_unique_rows = len(u_data)

# 1. 创建辅助张量A，并用一个大值（如1000，确保大于任何可能的行索引）初始化
# dtype应为long以匹配索引类型
placeholder_value = num_original_rows + 100 # 确保占位符大于最大行索引
A = placeholder_value * torch.ones((num_original_rows, num_unique_rows), dtype=torch.long)

# 2. 填充张量A
# A[i, inverse_indices[i]] = i
# torch.arange(num_original_rows) 生成 [0, 1, ..., num_original_rows-1]
# inverse_indices 提供了每个原始行对应的唯一行索引
# 这样，A[i, j] = i 当且仅当原始行 i 属于唯一行组 j
A[torch.arange(num_original_rows), inverse_indices] = torch.arange(num_original_rows)

# 3. 使用argmin查找首次出现索引
# 沿dim=0（列方向）查找最小值，即找到每个唯一行组的最小原始行索引
unique_indices_optimized = torch.argmin(A, dim=0)

print("优化方法得到的首次出现索引：", unique_indices_optimized)

# 验证两种方法结果是否一致
# （为了验证，这里重新计算了传统方法的结果）
unique_indices_traditional = torch.zeros(len(u_data), dtype=torch.long)
for idx in range(len(u_data)):
    unique_indices_traditional[idx] = torch.where(inverse_indices == idx)[0][0]

print("两种方法结果是否一致：", torch.allclose(unique_indices_optimized, unique_indices_traditional))

代码解释：

• placeholder_value = num_original_rows + 100: 我们选择一个肯定大于任何有效行索引（0到num_original_rows-1）的值作为占位符。
• A = placeholder_value * torch.ones(...): 初始化一个所有元素都是占位符值的二维张量A。
• A[torch.arange(num_original_rows), inverse_indices] = torch.arange(num_original_rows): 这是关键的向量化步骤。
  • torch.arange(num_original_rows) 生成一个从0到num_original_rows-1的序列，代表原始张量的行索引。
  • inverse_indices 包含了原始张量中每一行对应的唯一行索引。
  • 通过这种高级索引方式，我们将A中对应位置的值设置为原始行索引本身。例如，如果inverse_indices[5]是2，那么A[5, 2]将被设置为5。
• unique_indices_optimized = torch.argmin(A, dim=0): 对张量A的每一列（dim=0），argmin会返回最小值所在的行索引。由于有效值（原始行索引）都远小于占位符，并且这些值代表了原始行索引，argmin自然会找到属于该唯一行组的最小原始行索引，即首次出现的索引。

3. 效率与内存考量

• 效率提升： 优化方法消除了Python层面的显式循环和多次torch.where调用，转而使用高度优化的PyTorch张量操作（如高级索引和argmin），这在GPU上运行时尤其能体现出显著的性能优势。对于大规模数据，这种向量化处理通常比循环快几个数量级。
• 内存使用： 优化方法的主要缺点是它需要创建一个辅助的二维张量A，其大小为 (原始行数, 唯一行数)。如果原始张量行数和唯一行数都非常大，这个辅助张量可能会占用大量内存。例如，如果原始张量有100万行，其中有10万个唯一行，那么A将是 1,000,000 x 100,000 的张量，这可能导致内存溢出。

总结：

在选择方法时，需要根据实际应用场景进行权衡：

• 小到中等规模数据： 优化方法通常是更优的选择，因为它提供了显著的性能提升。
• 大规模数据且内存受限： 如果原始行数和唯一行数都非常庞大，以至于创建辅助张量A会导致内存问题，那么可能需要考虑其他更节省内存但可能效率稍低的方法，或者分块处理。

总而言之，通过巧妙地利用PyTorch的张量操作，我们可以将复杂的循环逻辑转化为高效的向量化计算，从而在处理数据时获得更好的性能。

以上就是PyTorch张量唯一行首次出现索引的高效查找方法的详细内容，更多请关注php中文网其它相关文章！