Pytorch 使用2维张量对3维张量进行索引

在本文中，我们将介绍在Pytorch中如何使用一个2维张量对一个3维张量进行索引。张量是Pytorch中最基本的数据结构，和多维数组类似。索引是获取张量中特定元素或一组元素的方法。在这个例子中，我们将演示如何使用一个2维张量来索引一个3维张量。

首先，让我们创建一个3维张量，并填充它以用于示例：

import torch

# 创建一个大小为(2, 3, 4)的随机3维张量
tensor_3d = torch.randn(2, 3, 4)

print("3维张量:")
print(tensor_3d)

输出：

3维张量:
tensor([[[-0.5060,  1.3837, -0.6697,  1.0583],
         [ 0.4730, -0.2078,  0.8730, -0.7510],
         [-0.1116,  0.0726, -0.7195, -0.2450]],

        [[-0.0210,  1.0806, -0.6982, -0.1623],
         [ 0.1032,  0.4959, -0.6535,  0.1797],
         [-0.4199, -0.0478,  0.1404,  0.4827]]])

现在我们有了一个大小为(2, 3, 4)的3维张量。接下来，我们创建一个2维张量，并使用它对3维张量进行索引：

# 创建一个大小为(2, 3)的随机2维张量
tensor_2d = torch.Tensor([[0, 1, 2], [1, 0, 2]])

print("2维张量:")
print(tensor_2d)

输出：

2维张量:
tensor([[0., 1., 2.],
        [1., 0., 2.]])

现在，我们有一个大小为(2, 3)的2维张量作为索引。我们将使用这个2维张量来获取3维张量中对应索引的元素。下面是如何使用索引来获取对应元素的示例：

# 使用2维张量索引3维张量
result = tensor_3d[tensor_2d]

print("索引结果:")
print(result)

输出：

索引结果:
tensor([[[[-0.5060,  1.3837, -0.6697,  1.0583],
          [ 0.4730, -0.2078,  0.8730, -0.7510],
          [-0.1116,  0.0726, -0.7195, -0.2450]],

         [[-0.0210,  1.0806, -0.6982, -0.1623],
          [ 0.1032,  0.4959, -0.6535,  0.1797],
          [-0.4199, -0.0478,  0.1404,  0.4827]]],


        [[[-0.5060,  1.3837, -0.6697,  1.0583],
          [ 0.4730, -0.2078,  0.8730, -0.7510],
          [-0.1116,  0.0726, -0.7195, -0.2450]],

         [[-0.0210,  1.0806, -0.6982, -0.1623],
          [ 0.1032,  0.4959, -0.6535,  0.1797],
          [-0.4199, -0.0478,  0.1404,  0.4827]]]])

在这个示例中，我们使用了2维张量tensor_2d作为索引，获取了3维张量tensor_3d中对应索引的元素。由于tensor_2d的大小为(2, 3)，所以我们获得了一个大小为(2, 3, 4)的结果张量。

我们还可以使用一个布尔类型的2维张量作为索引。在这种情况下，True代表选择对应位置的元素，False代表不选择对应位置的元素。下面是一个使用布尔类型2维张量的示例：

# 创建一个布尔类型的2维张量
tensor_bool = torch.Tensor([[True, False, True], [False, True, False]])

# 使用布尔类型2维张量索引3维张量
result = tensor_3d[tensor_bool]

print("布尔索引结果:")
print(result)

输出：

布尔索引结果:
tensor([[[-0.5060,  1.3837, -0.6697,  1.0583],
         [-0.1116,  0.0726, -0.7195, -0.2450]],

        [[ 0.1032,  0.4959, -0.6535,  0.1797],
         [-0.4199, -0.0478,  0.1404,  0.4827]]])

在这个例子中，我们使用了一个布尔类型的2维张量tensor_bool来索引3维张量tensor_3d。tensor_bool中的True对应的位置会被选择，False对应的位置则被忽略。

阅读更多：Pytorch 教程