PyTorch 구현 시 꼭 알아야 할 필수 함수 정리

PyTorch 구현 시 꼭 알아야 할 필수 함수 정리

PyTorch는 딥러닝 모델 구현에 강력하고 유연한 기능을 제공합니다.
이 문서는 모델 구현 시 자주 사용하는 핵심 함수들을 주제별로 정리한 참고용 자료입니다.

---

1. 텐서(Tensor) 생성 및 변형

기본 생성

torch.tensor(data)          # 파이썬 리스트나 넘파이 배열로부터 텐서 생성
torch.zeros(size)           # 0으로 채운 텐서
torch.ones(size)            # 1로 채운 텐서
torch.randn(size)           # 정규분포로 초기화된 텐서
torch.arange(start, end)    # 일정 간격으로 생성된 텐서
torch.linspace(start, end, steps)

변형 및 차원 조작

x.view(shape)               # 텐서 형태 변경 (메모리 공유)
x.reshape(shape)            # 형태 변경 (복사될 수도 있음)
x.unsqueeze(dim)            # 차원 추가 ([N] → [N, 1])
x.squeeze(dim)              # 크기 1인 차원 제거
x.permute(dims)             # 차원 순서 변경
x.transpose(dim0, dim1)     # 두 차원 교환
torch.cat([a, b], dim)      # 텐서 연결
torch.stack([a, b], dim)    # 새로운 차원으로 스택

2. 인덱싱 및 선택 함수

기본 인덱싱

x[index]                    # 일반 인덱싱
x[:, 0]                     # 특정 열 선택
x[mask]                     # 불리언 마스크 인덱싱

torch.gather(input, dim, index)
# 특정 인덱스 위치의 값을 모으기. 예: 어텐션 연산에서 사용

torch.index_select(input, dim, index)
# 특정 차원의 인덱스들만 추출

torch.take(input, index)
# 전체를 1차원으로 펼친 후 인덱스로 선택

torch.scatter(input, dim, index, src)
# gather의 반대. 지정된 위치에 값 배치

torch.masked_select(input, mask)
# 마스크된 요소만 선택

# gather_and_expand
def gather_and_expand(x, index, dim):
    out = torch.gather(x, dim, index)
    return out.expand_as(x)

3. 연산 및 통계

torch.sum(x, dim)
torch.mean(x, dim)
torch.std(x)
torch.var(x)
torch.max(x, dim)
torch.min(x, dim)
torch.argmax(x, dim)
torch.argmin(x, dim)
torch.clamp(x, min, max)
torch.norm(x)
torch.exp(x)
torch.log(x)
torch.softmax(x, dim)
torch.log_softmax(x, dim)
torch.sigmoid(x)
torch.relu(x)
torch.tanh(x)

4. 자동미분 (Autograd)

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
y.backward(torch.ones_like(x))  # dy/dx 계산
x.grad                           # 기울기 확인
x.detach()                       # 그래프 분리

5. 모델 구성 (nn 모듈)

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(128, 10)
    def forward(self, x):
        return self.fc(x)

# 자주 쓰는 레이어
nn.Linear(in_features, out_features)
nn.Conv2d(in_ch, out_ch, kernel_size)
nn.ReLU()
nn.Sigmoid()
nn.Tanh()
nn.BatchNorm1d(num_features)
nn.Dropout(p)
nn.Embedding(num_embeddings, embedding_dim)

6. 손실 함수 (Loss)

nn.MSELoss()
nn.CrossEntropyLoss()
nn.BCELoss()
nn.L1Loss()

7. 최적화 (Optimizer)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
optimizer.zero_grad()    # 기울기 초기화
loss.backward()          # 역전파
optimizer.step()         # 파라미터 업데이트

8. GPU 및 장치 설정

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = x.to(device)
model.to(device)

9. 유용한 헬퍼 함수

torch.no_grad()                         # 평가 시 gradient 비활성화
torch.nn.utils.clip_grad_norm_(params, max_norm)  # gradient 클리핑
torch.save(model.state_dict(), "model.pt")        # 모델 저장
model.load_state_dict(torch.load("model.pt"))     # 모델 로드