pytorch-patterns

PyTorch 개발 패턴

강력하고 효율적이며 재현 가능한 딥러닝 애플리케이션을 구축하기 위한 관용적인 PyTorch 패턴과 모범 사례입니다.

활성화 시점

• 새로운 PyTorch 모델 또는 학습 스크립트를 작성할 때
• 딥러닝 코드를 검토할 때
• 학습 루프 또는 데이터 파이프라인을 디버깅할 때
• GPU 메모리 사용량 또는 학습 속도를 최적화할 때
• 재현 가능한 실험을 설정할 때

핵심 원칙

1. 디바이스에 구애받지 않는 코드

디바이스를 하드코딩하지 않고 CPU와 GPU 모두에서 작동하는 코드를 항상 작성하세요.

# 좋음: 디바이스에 구애받지 않음
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
data = data.to(device)

# 나쁨: 하드코딩된 디바이스
model = MyModel().cuda()  # GPU가 없으면 충돌
model = MyModel().cuda()
data = data.cuda()

2. 재현성 우선

재현 가능한 결과를 위해 모든 난수 시드를 설정하세요.

# 좋음: 완전한 재현성 설정
def set_seed(seed: int = 42) -> None:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# 나쁨: 시드 제어 없음
model = MyModel()  # 실행할 때마다 가중치가 다름

3. 명시적인 형태 관리

항상 텐서 형태를 문서화하고 검증하세요.

# 좋음: 형태 주석이 달린 forward 패스
def forward(self, x: torch.Tensor) -> torch.Tensor:
    # x: (batch_size, channels, height, width)
    x = self.conv1(x)    # -> (batch_size, 32, H, W)
    x = self.pool(x)     # -> (batch_size, 32, H//2, W//2)
    x = x.view(x.size(0), -1)  # -> (batch_size, 32*H//2*W//2)
    return self.fc(x)    # -> (batch_size, num_classes)

# 나쁨: 형태 추적 없음
def forward(self, x):
    x = self.conv1(x)
    x = self.pool(x)
    x = x.view(x.size(0), -1)  # 이 크기는 무엇인가?
    return self.fc(x)           # 이게 작동할까?

모델 아키텍처 패턴

깔끔한 nn.Module 구조

# 좋음: 잘 구성된 모듈
class ImageClassifier(nn.Module):
    def __init__(self, num_classes: int, dropout: float = 0.5) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(64 * 16 * 16, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

# 나쁨: 모든 것이 forward 안에 있음
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = F.conv2d(x, weight=self.make_weight())  # 호출할 때마다 가중치 생성!
        return x

적절한 가중치 초기화

# 좋음: 명시적인 초기화
def _init_weights(self, module: nn.Module) -> None:
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
    elif isinstance(module, nn.BatchNorm2d):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)

model = MyModel()
model.apply(model._init_weights)

학습 루프 패턴

표준 학습 루프

# 좋음: 모범 사례를 포함한 완전한 학습 루프
def train_one_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
    scaler: torch.amp.GradScaler | None = None,
) -> float:
    model.train()  # 항상 학습 모드 설정
    total_loss = 0.0

    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad(set_to_none=True)  # zero_grad()보다 효율적

        # 혼합 정밀도 학습
        with torch.amp.autocast("cuda", enabled=scaler is not None):
            output = model(data)
            loss = criterion(output, target)

        if scaler is not None:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)

검증 루프

# 좋음: 적절한 평가
@torch.no_grad()  # torch.no_grad() 블록으로 감싸는 것보다 효율적
def evaluate(
    model: nn.Module,
    dataloader: DataLoader,
    criterion: nn.Module,
    device: torch.device,
) -> tuple[float, float]:
    model.eval()  # 항상 평가 모드 설정 — 드롭아웃 비활성화, 배치 정규화 실행 통계 사용
    total_loss = 0.0
    correct = 0
    total = 0

    for data, target in dataloader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        total_loss += criterion(output, target).item()
        correct += (output.argmax(1) == target).sum().item()
        total += target.size(0)

    return total_loss / len(dataloader), correct / total

데이터 파이프라인 패턴

사용자 정의 데이터셋

# 좋음: 타입 힌트가 포함된 깔끔한 데이터셋
class ImageDataset(Dataset):
    def __init__(
        self,
        image_dir: str,
        labels: dict[str, int],
        transform: transforms.Compose | None = None,
    ) -> None:
        self.image_paths = list(Path(image_dir).glob("*.jpg"))
        self.labels = labels
        self.transform = transform

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
        img = Image.open(self.image_paths[idx]).convert("RGB")
        label = self.labels[self.image_paths[idx].stem]

        if self.transform:
            img = self.transform(img)

        return img, label

효율적인 DataLoader 구성

# 좋음: 최적화된 DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,            # 학습을 위해 섞기
    num_workers=4,           # 병렬 데이터 로딩
    pin_memory=True,         # CPU->GPU 전송 속도 향상
    persistent_workers=True, # 에폭 간 워커 유지
    drop_last=True,          # BatchNorm을 위한 일관된 배치 크기
)

# 나쁨: 느린 기본값
dataloader = DataLoader(dataset, batch_size=32)  # num_workers=0, pin_memory 없음

가변 길이 데이터를 위한 사용자 정의 collate

# 좋음: collate_fn에서 시퀀스 패딩
def collate_fn(batch: list[tuple[torch.Tensor, int]]) -> tuple[torch.Tensor, torch.Tensor]:
    sequences, labels = zip(*batch)
    # 배치 내 최대 길이로 패딩
    padded = nn.utils.rnn.pad_sequence(sequences, batch_first=True, padding_value=0)
    return padded, torch.tensor(labels)

dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)

체크포인트 패턴

체크포인트 저장 및 로드

# 좋음: 모든 학습 상태를 포함한 완전한 체크포인트
def save_checkpoint(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    epoch: int,
    loss: float,
    path: str,
) -> None:
    torch.save({
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }, path)

def load_checkpoint(
    path: str,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
) -> dict:
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    model.load_state_dict(checkpoint["model_state_dict"])
    if optimizer:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint

# 나쁨: 모델 가중치만 저장 (학습 재개 불가)
torch.save(model.state_dict(), "model.pt")

성능 최적화

혼합 정밀도 학습

# 좋음: GradScaler를 사용한 AMP
scaler = torch.amp.GradScaler("cuda")
for data, target in dataloader:
    with torch.amp.autocast("cuda"):
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)

대규모 모델을 위한 그래디언트 체크포인팅

# 좋음: 메모리를 위해 연산 절충
from torch.utils.checkpoint import checkpoint

class LargeModel(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 역전파 시 활성화를 다시 계산하여 메모리 절약
        x = checkpoint(self.block1, x, use_reentrant=False)
        x = checkpoint(self.block2, x, use_reentrant=False)
        return self.head(x)

속도를 위한 torch.compile

# 좋음: 더 빠른 실행을 위해 모델 컴파일 (PyTorch 2.0+)
model = MyModel().to(device)
model = torch.compile(model, mode="reduce-overhead")

# 모드: "default" (안전), "reduce-overhead" (더 빠름), "max-autotune" (가장 빠름)

빠른 참조: PyTorch 관용구

관용구	설명
model.train() / model.eval()	학습/평가 전에 항상 모드 설정
torch.no_grad()	추론을 위해 그래디언트 비활성화
optimizer.zero_grad(set_to_none=True)	더 효율적인 그래디언트 클리어
.to(device)	디바이스에 구애받지 않는 텐서/모델 배치
torch.amp.autocast	2배 속도를 위한 혼합 정밀도
pin_memory=True	더 빠른 CPU→GPU 데이터 전송
torch.compile	속도를 위한 JIT 컴파일 (2.0+)
weights_only=True	안전한 모델 로딩
torch.manual_seed	재현 가능한 실험
gradient_checkpointing	메모리를 위해 연산 절충

피해야 할 안티 패턴

# 나쁨: 검증 중 model.eval()을 잊어버림
model.train()
with torch.no_grad():
    output = model(val_data)  # 드롭아웃이 아직 활성화! BatchNorm이 배치 통계를 사용!

# 좋음: 항상 평가 모드 설정
model.eval()
with torch.no_grad():
    output = model(val_data)

# 나쁨: 자동미분을 깨뜨리는 in-place 연산
x = F.relu(x, inplace=True)  # 그래디언트 계산을 깨뜨릴 수 있음
x += residual                  # In-place 덧셈이 autograd 그래프를 깨뜨림

# 좋음: Out-of-place 연산
x = F.relu(x)
x = x + residual

# 나쁨: 학습 루프 내에서 반복적으로 데이터를 GPU로 이동
for data, target in dataloader:
    model = model.cuda()  # 반복할 때마다 모델을 이동!

# 좋음: 루프 전에 한 번 모델 이동
model = model.to(device)
for data, target in dataloader:
    data, target = data.to(device), target.to(device)

# 나쁨: 역전파 전에 .item() 사용
loss = criterion(output, target).item()  # 그래프에서 분리!
loss.backward()  # 오류: .item()을 통해 역전파할 수 없음

# 좋음: 로깅에만 .item() 호출
loss = criterion(output, target)
loss.backward()
print(f"Loss: {loss.item():.4f}")  # 역전파 후 .item()은 괜찮음

# 나쁨: torch.save를 제대로 사용하지 않음
torch.save(model, "model.pt")  # 전체 모델 저장 (취약, 이식 불가)

# 좋음: state_dict 저장
torch.save(model.state_dict(), "model.pt")

기억하세요: PyTorch 코드는 디바이스에 구애받지 않고, 재현 가능하며, 메모리를 의식해야 합니다. 의심스러우면 torch.profiler로 프로파일링하고 torch.cuda.memory_summary()로 GPU 메모리를 확인하세요.