pytorch索引操作函数介绍

PyTorch 提供了一系列强大的索引操作函数，用于高效地操作张量数据。以下是常用的 PyTorch 索引操作函数及其用途：

1. 基础索引操作

这些操作类似于 NumPy 的基本索引方式。

1.1 方括号索引 ([])

• 支持 切片、布尔索引、高级索引。
• 示例：

x = torch.tensor([[10, 20, 30], [40, 50, 60]])
print(x[0, 1])  # 取第 0 行第 1 列，输出 20
print(x[:, 1])  # 取所有行的第 1 列，输出 tensor([20, 50])
print(x[1])     # 取第 1 行，输出 tensor([40, 50, 60])

2. 索引操作函数

2.1 torch.index_select

按指定维度的索引提取元素。

• 用法：

torch.index_select(input, dim, index, *, out=None)

• 参数：
  • input：输入张量。
  • dim：指定的维度。
  • index：索引张量，必须是一维张量。
• 示例：

x = torch.tensor([[10, 20, 30], [40, 50, 60]])
index = torch.tensor([0, 2])
result = torch.index_select(x, dim=1, index=index)
print(result)  # 输出 tensor([[10, 30], [40, 60]])

2.2 torch.gather

根据索引张量，从输入张量中收集数据。

• 用法：

torch.gather(input, dim, index, *, sparse_grad=False, out=None)

• 参数：
  • input：输入张量。
  • dim：指定的维度。
  • index：索引张量，必须与输入张量形状兼容。
• 示例：

x = torch.tensor([[10, 20, 30], [40, 50, 60]])
index = torch.tensor([[0, 2], [1, 0]])
result = torch.gather(x, dim=1, index=index)
print(result)  # 输出 tensor([[10, 30], [50, 40]])

2.3 torch.scatter

将源张量的值按索引写入目标张量。

• 用法：

torch.scatter(input, dim, index, src, *, reduce=None)

示例：

x = torch.zeros(2, 3)
index = torch.tensor([[0, 1, 2], [0, 1, 2]])
src = torch.tensor([[10., 20., 30.], [40., 50., 60.]])
result = torch.scatter(x, dim=1, index=index, src=src)
print(result)  # 输出 tensor([[10., 20., 30.], [40., 50., 60.]])

2.4 torch.take

按照扁平化索引从张量中提取元素。

• 用法：

torch.take(input, index)

示例：

x = torch.tensor([[10, 20, 30], [40, 50, 60]])
index = torch.tensor([0, 2, 5])
result = torch.take(x, index)
print(result)  # 输出 tensor([10, 30, 60])

torch.gather和 torch.scatter区别

2.5 torch.masked_select

根据布尔掩码从张量中提取元素。

• 用法：

torch.masked_select(input, mask, *, out=None)

示例：

x = torch.tensor([[10, 20, 30], [40, 50, 60]])
mask = x > 30
result = torch.masked_select(x, mask)
print(result)  # 输出 tensor([40, 50, 60])

2.6 torch.nonzero

返回所有非零元素的索引。

• 用法：

torch.nonzero(input, *, as_tuple=False)

示例：

x = torch.tensor([[0, 1, 0], [2, 0, 3]])
result = torch.nonzero(x)
print(result)  # 输出 tensor([[0, 1], [1, 0], [1, 2]])

2.7 torch.where

根据条件选择值。

• 用法：

torch.where(condition, x, y)

示例：

x = torch.tensor([1, 2, 3])
y = torch.tensor([10, 20, 30])
condition = x > 1
result = torch.where(condition, x, y)
print(result)  # 输出 tensor([10,  2,  3])

2.8 torch.advanced_indexing

• PyTorch 支持 布尔张量索引 和 整形张量索引。
• 示例：

x = torch.tensor([[10, 20, 30], [40, 50, 60]])
index = torch.tensor([[0, 1], [1, 0]])
result = x[index]
print(result)

输出：

tensor([[[10, 20, 30],
         [40, 50, 60]],

        [[40, 50, 60],
         [10, 20, 30]]])

注：第0维度操作

3. 总结

这些索引操作函数可以帮助简化张量操作，提高代码效率。