выбрать подэлементы из другой партии

Question

У меня есть первый тензор размером torch.Size([12, 64, 8, 8, 3]) , (8,8,3) - размер изображения, 64 - патчи, 12 - размер пакета

Есть другой тензор с размером torch.Size([12, 10]) которые выбирают 10 патчей для каждого элемента в пакете (выберите 10 патчей из 64). Таким образом, он хранит индексы. Как использовать это для запроса первого тензора с пониманием списка?

Answer 1

Цель состоит в том, чтобы запросить 1-й патч (из 10 выбранных патчей) каждого пакета. При итерации по b получаем список индексов выбранных патчей. Выберите 1-й из них по индексу 0. Так как они являются тензорами, преобразуйте тип в int, чтобы индексировать в тензор изображений и получить соответствующий патч для каждого пакета.

a = torch.rand(12, 64, 8, 8, 3) # generating 12 batches, with 64 patches,each of size 8x8x3
b = torch.randint(64, (12, 10)) # choosing 10 patches (within the 64), for each of the 12 batches

first_tensors = [a[batch, int(patches[0])] for batch, patches in zip(range(12), b)]

Для для ясности, понимание нижеследующего списка даст индексы 1-го патча каждой партии.

[[batch, int(patches[0])] for batch, patches in zip(range(12), b)]
[[0, 40],
 [1, 27],
 [2, 17],
 [3, 62],
 [4, 9],
 [5, 51],
 [6, 32],
 [7, 38],
 [8, 63],
 [9, 10],
 [10, 2],
 [11, 6]]

Индексирование тензора изображений a с каждой парой индексов в приведенном выше списке даст соответствующие патч.

Answer 2

Вы можете использовать index_select:

c = [torch.index_select(i, dim=0, index=j) for i, j in zip(a,b)]

a и b - ваш тензор и индексы соответственно.

Впоследствии вы можете stack это в нулевом измерении .