Два аргумента этой функции, index и dim, являются ключом к пониманию функции.
В случае 2D dim = 0 соответствует строкам, а dim = 1 соответствует столбцам.
В случае 3D dim = 0 соответствует изображению из пакета, dim = 1 соответствует строкам, а dim = 2 соответствует столбцам.
Случай двумерного входного тензора
1. Понимание аргумента dim:
а. Когда dim = 0, мы выбираем строки.
б. Когда dim = 1, мы выбираем столбцы.
2. Понимание аргумента индекса:
а. Аргумент индекса будет иметь такое же количество измерений, что и входные данные (не означает, что форма будет такой же).
б. Выходной тензор будет иметь ту же форму, что и индексный тензор.
в. Элементы индексного тензора сообщают, какую строку (для dim = 0, двумерный случай) выбрать, а положение конкретного элемента указывает, какой столбец выбрать.
д. Элементы индексного тензора говорят, какой столбец (для dim = 1, двумерный случай) выбрать, а положение конкретного элемента говорит, какую строку выбрать.
Случай трехмерного входного тензора
1. Понимание аргумента dim:
а. Когда dim = 0, мы выбираем изображение из пакета.
б. Когда dim = 1, мы выбираем строки.
в. Когда dim = 2, мы выбираем столбцы.
2. Понимание аргумента индекса:
а. a и b сверху, применяется случай двумерного входного тензора.
б. Элементы индексного тензора сообщают, какое изображение из пакета (для случая dim = 0, 3D) выбрать, а положение конкретного элемента указывает, какие строки и столбцы выбрать, и так далее для dim = 1 и dim = 2.
Давайте возьмем два примера для случая 2D.
1-й пример
Когда тусклый = 0, и
ind_2d = [[3, 2, 0, 1]]
ind_2d имеет форму (1, 4), поэтому выходные данные будут иметь такую же форму.
0-й элемент ind_2d, т. е. 3, говорит, что мы выбираем 3-ю строку и 0-й столбец (поскольку 3 — это 0-й элемент индексного тензора).
1-й элемент ind_2d, то есть 2, говорит, что мы выбираем 2-ю строку, начиная с 0 (строка, потому что dim = 0) и 1-й столбец (начиная с 1-го элемента индексного тензора). И так далее.
2-й пример
Предположим, что из приведенного выше тензора src_2d мы хотим выбрать 0, 6, 10 и 15. (Совет: мы читаем эти числа по порядку сверху вниз, поэтому мы формируем столбец, подобный тензору индексов. Когда мы читаем слева направо, мы формируем строку как индексный тензор.)
Теперь 0 принадлежит 0-му столбцу, 6 принадлежит 2-му столбцу, 10 принадлежит 2-му столбцу и 15 принадлежит 3-му столбцу.
Таким образом, наш индексный тензор имеет вид [[0, 2, 2, 3]] в виде столбца, т. е. формы (4, 1). А так как мы выбрали столбцы, то dim = 1.
Мы никак не могли выбрать строки и получить желаемый выходной тензор.
Случай с 3D очень похож. Получите идеи из инфографики ниже.
Вот ссылка на блокнот Python с несколькими примерами для двухмерного и трехмерного входного тензора:
Ссылка на функцию Pytorch:
https://pytorch.org/docs/stable/generated/torch.gather.html?highlight=torch%20gather#torch.gather
