import torch
a = torch.randn(3,5)
print(a)
# 下行会有错误 IndexError: shape mismatch: indexing tensors could not be broadcast together with shapes [2], [3]
#b = a[[0,2],[1,3,4]]
# 改成
b = a[[0,2],:][:,[1,3,4]]
print(b)
输出是:
tensor([[ 0.3627, -0.7073, -0.3986, 2.0124, 0.4630],
[ 0.2778, -0.3370, -0.6217, -1.0586, -0.1732],
[-1.9196, 0.6089, -0.3361, -0.2576, 1.2209]])
tensor([[-0.7073, 2.0124, 0.4630],
[ 0.6089, -0.2576, 1.2209]])
结果是取了第0、2行、第1、3、4列的子矩阵。
有可能这不是想要的结果,比如如下代码:
import torch
a = torch.randn(3,5)
b = a[[0,2],[1,3]]
c = a[[0,2],:][:,[1,3]]
print(a)
print(b)
print(c)
输出是:
tensor([[-0.9547, -0.2727, -0.5674, -1.2124, -0.4707],
[ 1.0657, -0.9457, 0.3516, 0.5730, -0.2926],
[-0.7380, 0.6028, 0.2395, 0.2386, 0.9953]])
tensor([-0.2727, 0.2386])
tensor([[-0.2727, -1.2124],
[ 0.6028, 0.2386]])
b和c是完全不同的。