인덱싱과 슬라이싱

NumPy 인덱싱에 몇 가지만 더하면 돼

NumPy 인덱싱에 익숙하다면 다음 10분은 복습에 가까워. PyTorch의 규칙은 이래:

기본 인덱싱: 정수와 슬라이스를 사용하며, 가능하면 복사하지 않고 뷰를 반환해.
불리언 인덱싱: 모양이 맞는 bool 텐서로 원소를 골라. 선택한 원소를 담은 1D 텐서를 새로 만들어.
고급 인덱싱: 정수 텐서를 인덱스로 사용하며 항상 복사본을 만들어.
혼합 인덱싱: 위 방식을 조합할 수 있고 규칙은 NumPy와 같아.

가장 중요한 사실은 기본 슬라이싱이 복사본이 아니라 뷰를 반환한다는 거야. 슬라이스에 값을 쓰면 원본도 바뀌어. 제자리 갱신을 저렴하게 만드는 의도된 기능이지만, '텐서 하나만 바꿨는데 셋이 함께 변했다'는 버그의 원인이 되기도 해.

자주 쓰는 선택 방식

x[:, 0]: 모든 행에서 첫 번째 열을 골라.
x[..., -1]: x의 차원 수와 상관없이 마지막 차원의 마지막 원소를 골라. 짧고 읽기 좋아.
x[None, :] / x.unsqueeze(0): 배치 차원을 추가해.
x[mask]: 불리언 마스크로 걸러.
torch.gather(x, dim, index): 인덱스 텐서를 따라 특정 차원의 값을 골라. 어텐션과 손실 계산에서 아주 자주 써.

Code

기본 인덱싱과 슬라이싱·python

import torch

t = torch.tensor([
    [ 1,  2,  3,  4],
    [ 5,  6,  7,  8],
    [ 9, 10, 11, 12],
])

t[0]          # tensor([1, 2, 3, 4])  — first row
t[0, 2]       # tensor(3)             — single element
t[-1]         # last row
t[:, 1]       # all rows, col 1 → tensor([2, 6, 10])
t[0:2, :]     # first two rows
t[::2, :]     # every other row
t[..., -1]    # last col, regardless of rank → tensor([4, 8, 12])

불리언 인덱싱과 고급 인덱싱·python

import torch

t = torch.tensor([10, 20, 30, 40, 50])

# Boolean — same-shape bool tensor
mask = t > 25
t[mask]                  # tensor([30, 40, 50])

# Fancy — integer index tensor
idx = torch.tensor([0, 2, 4])
t[idx]                   # tensor([10, 30, 50])

# Mix in 2D
m = torch.arange(20).reshape(4, 5)
m[m % 3 == 0]            # all multiples of 3, flattened
m[[0, 2, 3]]             # rows 0, 2, 3
m[[0, 2, 3], [1, 1, 1]]  # element-wise: (0,1), (2,1), (3,1)

뷰 대 복사: 함정·python

import torch

a = torch.arange(12).reshape(3, 4)
row = a[1]               # VIEW into a (no copy)
row[:] = 99              # writes into a too!
print(a)
# tensor([[ 0,  1,  2,  3],
#         [99, 99, 99, 99],
#         [ 8,  9, 10, 11]])

# To get an independent copy:
row_copy = a[1].clone()
row_copy[:] = -1
print(a[1])              # untouched: tensor([99, 99, 99, 99])

torch.gather: 어텐션과 손실 계산의 일꾼·python

import torch

# Pick the predicted-class probability for each sample
logits = torch.randn(4, 5).softmax(-1)   # (batch=4, classes=5)
labels = torch.tensor([2, 0, 4, 1])

# We want logits[i, labels[i]] for each i
chosen = logits.gather(1, labels.unsqueeze(1)).squeeze(1)
print(chosen.shape)  # torch.Size([4])

# Equivalent (less efficient) loop:
# chosen = torch.tensor([logits[i, labels[i]] for i in range(4)])

Exercise

torch.randn으로 (4, 5) 로짓 텐서를 만들어. 레이블 텐서 [2, 0, 4, 1]이 주어졌을 때 각 샘플의 정답 클래스 로짓을 Python 반복문과 torch.gather 두 방식으로 추출해. 값이 같은지 확인하고, timeit으로 (1024, 1000) 텐서에서 속도를 비교해. torch.gather가 훨씬 빨라야 해.

NumPy 인덱싱에 몇 가지만 더하면 돼

자주 쓰는 선택 방식

Code

External links

Exercise

Progress

댓글 0