理解PyTorch 张量的多维张量索引

发布于:2025-03-20 ⋅ 阅读:(14) ⋅ 点赞:(0)

PyTorch 的 多维张量索引 非常强大,但理解索引的取值有一定的难度。下面以具体的示例数据解释这种复杂的索引操作。

1.  单个多维张量索引

示例代码:
import numpy as np

data = torch.arange(12).view(3, 4)
print(f"输入数据data为:\n{data},\n形状为:{data.shape}")
print("--"*10)

## 一维张量
one_dim_index = torch.tensor([0, 1])
print(f"一维张量为:\n{one_dim_index},\n形状为:{one_dim_index.shape}\n")
result = data[one_dim_index]
print(f"取索引后数据为:\n{result},\n形状为:{result.shape}")
print("--"*10)

## 创建不同形式的多维索引
multi_dim_index1 = torch.tensor([[0], [1]])
print(f"多维张量为:\n{multi_dim_index1},\n形状为:{multi_dim_index1.shape}\n")
result1 = data[multi_dim_index1]
print(f"取索引后数据为:\n{result1},\n形状为:{result1.shape}")
print("--"*10)

multi_dim_index2 = torch.tensor([[[0], [1]]])
print(f"多维张量为:\n{multi_dim_index2},\n形状为:{multi_dim_index2.shape}\n")
result2 = data[multi_dim_index2]
print(f"取索引后数据为:\n{result2},\n形状为:{result2.shape}")
print("--"*10)

multi_dim_index3 = torch.tensor([[[0, 1]]])
print(f"多维张量为:\n{multi_dim_index3},\n形状为:{multi_dim_index3.shape}\n")
result3 = data[multi_dim_index3]
print(f"取索引后数据为:\n{result3},\n形状为:{result3.shape}")
print("--"*10)
代码逐行解读
1. 导入必要的库