Pytorch: разрезание 3D-тензора по измерению 2 с использованием 1D-тензора в качестве индекса ⇐ Python

[Anonymous]

Я хочу разрезать 3D-тензор в PyTorch. Форма 3D-тензора src_tensor равна (batch, max_len, Hidden_dim), и у меня есть индексные векторные индексы 1D с формой (batch,). Я хочу разрезать второе измерение src_tensor. Я могу добиться этой функциональности с помощью следующих кодов:

Код: Выделить всё

import torch
nums = 30
l = [i for i in range(nums)]
src_tensor = torch.Tensor(l).reshape((3,5,2))
indices = [1,2,3]
slice_tensor = torch.zeros((3,2,2))
for i in range(3):
p1,p2 = indices[i],indices[i]+1
slice_tensor[i,:,:]=src_tensor[i,[p1,p2],:]
print(src_tensor)
print(indices)
print(slice_tensor)
"""
tensor([[[ 0.,  1.],
[ 2.,  3.],
[ 4.,  5.],
[ 6.,  7.],
[ 8.,  9.]],

[[10., 11.],
[12., 13.],
[14., 15.],
[16., 17.],
[18., 19.]],

[[20., 21.],
[22., 23.],
[24., 25.],
[26., 27.],
[28., 29.]]])
[1, 2, 3]
tensor([[[ 2.,  3.],
[ 4.,  5.]],

[[14., 15.],
[16., 17.]],

[[26., 27.],
[28., 29.]]])
"""

Мой вопрос заключается в том, можно ли упростить приведенный выше код, например, исключив цикл for.


Подробнее здесь: https://stackoverflow.com/questions/791 ... -the-index