PyTorch张量索引用法速查

作为数据科学家或软件工程师，你可能经常处理大型数据集和复杂的数学运算，这些运算需要高效且可扩展的计算。PyTorch 是一个流行的开源机器学习库，它通过 GPU 加速提供快速灵活的张量计算。在本文中，我们将深入研究 PyTorch 张量索引，这是一种强大的技术，可让你轻松选择和操作张量的特定元素或子集。

NSDT工具推荐： Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割

1、什么是 PyTorch 张量索引？

张量索引是根据张量的位置或值选择特定元素或子集的过程。PyTorch 张量索引提供了一组丰富的索引操作，使你可以使用不同的索引方案（例如整数索引、布尔索引和高级索引）选择和修改张量元素。

在 PyTorch 中，张量是一个多维数组，可以存储不同类型和大小的数值数据。可以使用一个或多个索引对张量进行索引，这些索引指定元素沿张量每个维度的位置。索引张量将返回一个包含所选元素的新张量或具有修改元素的原始张量的视图。

1、整数索引

整数索引是张量索引的最基本形式，它允许你使用张量沿每个维度的整数位置来选择特定元素。你可以使用整数索引来选择单个元素或子张量，方法是提供与所需元素的索引相对应的整数列表。

import torch

# create a 2D tensor of size (3, 4)
x = torch.tensor([[1, 2, 3, 4],
                  [5, 6, 7, 8],
                  [9, 10, 11, 12]])

# select the element at position (1, 2)
print(x[1, 2])   # output: tensor(7)

# select the sub-tensor of size (2, 3) starting at position (0, 1)
print(x[0:2, 1:4])   # output: tensor([[2, 3, 4], [6, 7, 8]])

2、布尔索引

布尔索引允许你根据布尔条件选择张量的特定元素。你可以使用布尔索引来选择满足特定条件的元素或屏蔽不满足条件的元素。

# create a 1D tensor of size 5
x = torch.tensor([1, 2, 3, 4, 5])

# select the elements that are greater than 3
print(x[x > 3])   # output: tensor([4, 5])

# mask out the elements that are not even
x[x % 2 != 0] = 0
print(x)   # output: tensor([0, 2, 0, 4, 0])

3、高级索引

高级索引允许你使用整数索引、布尔索引和其他高级索引方案的组合来选择张量的特定元素。你可以使用高级索引从张量的多个维度中选择元素，或创建具有不同形状的新张量。

import torch

# create a 3D tensor of size (2, 3, 4)
x = torch.tensor([[[1, 2, 3, 4],
                   [5, 6, 7, 8],
                   [9, 10, 11, 12]],
                  
                  [[13, 14, 15, 16],
                   [17, 18, 19, 20],
                   [21, 22, 23, 24]]])

# select the elements at positions (0, 1, 2) and (1, 2, 3)
print(x[:, (0, 1, 2), (1, 2, 3)])   # output: tensor([[ 2,  7, 12],
                                    #          [14, 19, 24]])

# create a new tensor by selecting the elements that are greater than 10
y = x[x > 10]
print(y)   # output: tensor([11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24])

4、修改张量元素

PyTorch 张量索引不仅允许你选择张量元素，还允许你就地修改它们或使用修改后的元素创建新的张量。你可以使用索引为张量元素分配新值或对选定元素应用数学运算。

# create a 2D tensor of size (3, 3)
x = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

# assign a new value to the element at position (1, 2)
x[1, 2] = 10
print(x)   # output: tensor([[ 1,  2,  3], [ 4,  5, 10], [ 7,  8,  9]])

# multiply the elements that are greater than 5 by 2
x[x > 5] *= 2
print(x)   # output: tensor([[ 1,  2,  3], [ 4,  5, 20], [14, 16, 18]])