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)