Pytorch - gather

이번 포스트에서는 Pytorch gather 함수를 알아보도록 하겠습니다. gather 함수는 scatter 함수와 정반대로 동작하는 함수로 주어진 "input" 텐서에서 "dim", "index" 에 맞게 값을 모아주는 역할을 합니다.

torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor

gather 함수는 "input" 텐서에서 "dim", "index" 에 해당하는 요소값으로 새로운 텐서를 구성합니다. 3차원 텐서에 대해서 다음과 같이 동작합니다. scatter 함수와 반대 동작인 것을 알 수 있습니다.

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

정의된 파라미터는 다음과 같습니다.

• input (Tensor): 입력 텐서
• dim (int): 입력 텐서의 기준 축
• index (LongTensor): 입력 텐서를 모아줄 인덱스로 구성된 텐서로 정수형 텐서이어야 합니다.
• out(Tensor, optional): gather 함수를 담을 텐서로 따로 선언할 경우 in-place 함수로 동작합니다.

Example

다음 예를 통해 동작을 확인해 보겠습니다.

t = torch.tensor([[1, 2], [3, 4]])
torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]]))
# tensor([[ 1,  1],
#         [ 4,  3]])

scatter 함수와 마찬가지로 "index" 에 주어진 텐서 기준으로 동작합니다. 특히, "index" 텐서의 구성요소 별로 새로운 텐서 "out" 을 구성하다보니 "out" 텐서와 "index" 텐서는 크기가 같습니다.

• out[0][0] 에는 t[0][index[0][0]] = t[0][0] = 1 이 대입됩니다.
• out[0][1] 에는 t[0][index[0][1]] = t[0][0] = 1 이 대입됩니다.
• out[1][0] 에는 t[1][index[1][0]] = t[1][1] = 4 가 대입됩니다.
• out[1][1] 에는 t[1][index[1][1]] = t[1][0] = 3 이 대입됩니다.

"input" 과 "index" 는 같은 차원을 가져야 하며 차원 $d\neq dim$에 대해 $index.size(d)\leq input.size(d)$를 만족해야 합니다. ("dim" 이 아닌 차원에 대해 "index" 에 해당하는 행,열이 "input" 에도 존재해야 하기 때문입니다.) 또한, "out" 아규먼트에 어떠한 크기의 텐서를 넣어도 gather 함수의 결과가 할당됩니다. ("index" 와 같은 크기일 필요가 없습니다.)

 

참조

• https://pytorch.org/docs/stable/generated/torch.gather.html

 

---

Pytorch - scatter