고급 인덱싱

정수 배열 인덱싱 (Fancy Indexing)

정수 인덱스 목록(리스트 또는 텐서)을 사용해 임의 순서 로 원소를 선택합니다. 결과는 인덱스 배열과 같은 shape를 가집니다.

import torch

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

# 정수 리스트로 인덱싱
idx = [2, 0, 4, 1]
print(v[idx])  # tensor([30, 10, 50, 20])

# 중복 인덱스 허용
print(v[[0, 0, 2, 2]])  # tensor([10, 10, 30, 30])

# torch.tensor로 인덱싱
idx_t = torch.tensor([4, 3, 2, 1, 0])
print(v[idx_t])  # tensor([50, 40, 30, 20, 10])  ← 역순

2D에서의 Fancy Indexing

import torch

m = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

# 행 선택 (순서 변경 가능)
print(m[[2, 0, 1]])
# tensor([[7, 8, 9],
#         [1, 2, 3],
#         [4, 5, 6]])

# 행과 열을 동시에 (쌍으로 매칭)
rows = [0, 1, 2]
cols = [0, 1, 2]
print(m[rows, cols])  # tensor([1, 5, 9])  ← 대각선 원소

불리언 마스크 인덱싱

True/False 로 구성된 마스크 텐서로 원소를 선택합니다. 마스크와 텐서의 shape가 같아야 합니다.

import torch

v = torch.tensor([1.0, -2.0, 3.0, -4.0, 5.0])

# 양수만 선택
mask = v > 0
print(mask)   # tensor([ True, False,  True, False,  True])
print(v[mask])  # tensor([1., 3., 5.])

# 절댓값이 2 이상인 원소
mask2 = v.abs() >= 2
print(v[mask2])  # tensor([-2.,  3., -4.,  5.])

2D 마스크 인덱싱

import torch

m = torch.tensor([[1, -2, 3],
                  [-4, 5, -6],
                  [7, -8, 9]])

# 음수 원소만 0으로 교체
mask = m < 0
m_clean = m.clone()
m_clean[mask] = 0
print(m_clean)
# tensor([[1, 0, 3],
#         [0, 5, 0],
#         [7, 0, 9]])

# 마스크로 선택한 결과는 항상 1D
print(m[mask])   # tensor([-2, -4, -6, -8])
print(m[mask].shape)  # torch.Size([4])

조건 조합

import torch

x = torch.randn(100)

# AND 조건: 0 이상이고 1 미만
mask = (x >= 0) & (x < 1)

# OR 조건: 2보다 크거나 -2보다 작음
outlier_mask = (x > 2) | (x < -2)

# NOT 조건
not_mask = ~mask

print(f"0~1 사이 원소 수: {mask.sum().item()}")
print(f"이상치 수: {outlier_mask.sum().item()}")

torch.where() — 조건부 선택

torch.where(condition, x, y) 는 조건이 True 인 위치에는 x 값을, False 인 위치에는 y 값을 선택합니다. 결과 shape는 입력과 같습니다.

import torch

a = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
b = torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0])
condition = torch.tensor([True, False, True, False, True])

result = torch.where(condition, a, b)
print(result)  # tensor([ 1., 20.,  3., 40.,  5.])
#                          a   b    a   b    a

스칼라와 함께 사용

import torch

x = torch.tensor([-3.0, -1.0, 0.0, 1.0, 3.0])

# 양수는 그대로, 음수는 0으로 (ReLU 구현)
relu = torch.where(x > 0, x, torch.zeros_like(x))
print(relu)  # tensor([0., 0., 0., 1., 3.])

# 조건에 따라 스칼라 값 대입
clamped = torch.where(x > 2, torch.tensor(2.0), x)
print(clamped)  # tensor([-3., -1.,  0.,  1.,  2.])

인덱스 위치 반환: torch.where(condition)

조건만 전달하면 True 인 위치의 인덱스를 반환합니다.

import torch

x = torch.tensor([1.0, -2.0, 3.0, -4.0, 5.0])
indices = torch.where(x > 0)
print(indices)        # (tensor([0, 2, 4]),)  ← 튜플 형태
print(indices[0])     # tensor([0, 2, 4])

# 2D에서 행, 열 인덱스 반환
m = torch.tensor([[1, -2], [-3, 4]])
rows, cols = torch.where(m > 0)
print(rows)   # tensor([0, 1])
print(cols)   # tensor([0, 1])

실전 예시: 임계값 필터링

import torch

# 예측 확률 텐서 (배치 크기 = 8)
probs = torch.tensor([0.91, 0.43, 0.78, 0.12, 0.95, 0.55, 0.33, 0.87])

# 방법 1: 불리언 마스크 — 값 추출
threshold = 0.5
high_conf = probs[probs >= threshold]
print(high_conf)   # tensor([0.91, 0.78, 0.95, 0.55, 0.87])

# 방법 2: torch.where — 낮은 확률을 0으로 마스킹
masked = torch.where(probs >= threshold, probs, torch.zeros_like(probs))
print(masked)
# tensor([0.91, 0.00, 0.78, 0.00, 0.95, 0.55, 0.00, 0.87])

# 방법 3: 높은 신뢰도 인덱스 추출
high_idx = (probs >= threshold).nonzero(as_tuple=True)[0]
print(high_idx)    # tensor([0, 2, 4, 5, 7])
print(probs[high_idx])  # tensor([0.91, 0.78, 0.95, 0.55, 0.87])

기법	반환 형태	특징
`tensor[mask]`	1D (선택된 값)	선택된 원소만 추출
`torch.where(c, x, y)`	원본과 같은 shape	조건부 값 교체
`torch.where(c)`	인덱스 튜플	위치 정보 필요 시
`.nonzero()`	(N, ndim) 인덱스	0이 아닌 원소 위치