PyTorch에서 2 차원 이상의 인덱스를 선택하는 방법은 무엇입니까?

Question

가 a[2, :, 0, :] 그러나

0.5580 -0.0337 1.0048 -0.5044 0.6784 -1.6117 1.0084 1.1886 0.1278 0.3739 [torch.FloatTensor of size 2x5]

제공하는 동안 내가 ((2, :, 0, :), (1, :, 1, :), (2, :, 2, :), (0, :, 3, :) 같은 크기 (2, 4, 5) 또는 (4, 2, 5)?

의 결과 텐서를 하위 텐서를 선택하는 방법, a = torch.randn(3, 2, 4, 5)을 감안할 때, a[[2, 1, 2, 0], :, [0, 1, 2, 3], :]는

TypeError: Performing basic indexing on a tensor and encountered an error indexing dim 0 with an object of type list. The only supported types are integers, slices, numpy scalars, or if indexing with a torch.LongTensor or torch.ByteTensor only a single Tensor may be passed.

제공 numpy은 (4, 2, 5)을 반환하지만 성공적으로 ensor.

Answer 1

효과가 있습니까?

import torch 

a = torch.randn(3, 2, 4, 5) 
print(a.size()) 

b = [a[2, :, 0, :], a[1, :, 1, :], a[2, :, 2, :], a[0, :, 3, :]] 
b = torch.stack(b, 0) 

print(b.size()) # torch.Size([4, 2, 5])