PyTorch

실험 추적, 데이터셋 버전 관리, 메트릭·그라디언트·모델 로깅을 위해 W&B를 PyTorch와 통합하세요.

머신 러닝 실험 추적, 데이터셋 버전 관리, 프로젝트 협업에 W&B를 사용하세요.

​

이 노트북에서 다루는 내용

PyTorch 코드에 W&B를 통합해 파이프라인에 실험 추적을 추가하는 방법을 설명합니다.

# 라이브러리를 임포트합니다
import wandb

# config에 하이퍼파라미터 딕셔너리를 저장합니다
config = {
    "learning_rate": 0.001,
    "epochs": 100,
    "batch_size": 128
}

# 새 실험을 시작합니다
with wandb.init(project="new-sota-model", config=config) as run:

    # 모델과 데이터를 준비합니다
    model, dataloader = get_model(), get_data()

    # 선택 사항: 그라디언트를 추적합니다
    run.watch(model)

    for batch in dataloader:
    metrics = model.training_step()
    # 트레이닝 루프 내부에서 메트릭을 로깅해 모델 성능을 시각화합니다
    run.log(metrics)

    # 선택 사항: 마지막에 모델을 저장합니다
    model.to_onnx()
    run.save("model.onnx")

비디오 튜토리얼을 보며 따라 해보세요.

참고: Step으로 시작하는 섹션만 보면 기존 파이프라인에 W&B를 통합하는 데 필요한 내용은 충분합니다. 나머지는 데이터를 로드하고 모델을 정의하는 내용입니다.

​

설치, 임포트, 로그인

import os
import random

import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from tqdm.auto import tqdm

# 결정론적 동작 보장
torch.backends.cudnn.deterministic = True
random.seed(hash("setting random seeds") % 2**32 - 1)
np.random.seed(hash("improves reproducibility") % 2**32 - 1)
torch.manual_seed(hash("by removing stochasticity") % 2**32 - 1)
torch.cuda.manual_seed_all(hash("so runs are repeatable") % 2**32 - 1)

# 디바이스 설정
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# MNIST 미러 목록에서 느린 미러 제거
torchvision.datasets.MNIST.mirrors = [mirror for mirror in torchvision.datasets.MNIST.mirrors
                                      if not mirror.startswith("http://yann.lecun.com")]

​

Step 0: W&B 설치

시작하려면 먼저 라이브러리를 설치해야 합니다. wandb는 pip를 사용해 쉽게 설치할 수 있습니다.

!pip install wandb onnx -Uq

​

Step 1: W&B 임포트 및 로그인

웹 서비스에 데이터를 로깅하려면 로그인해야 합니다.

W&B를 처음 사용하는 경우 표시되는 링크를 통해 무료 계정에 가입해야 합니다.

import wandb

wandb.login()

​

실험과 파이프라인 정의하기

​

wandb.init()으로 메타데이터와 하이퍼파라미터 추적하기

프로그래밍 방식으로는 먼저 실험을 정의합니다. 하이퍼파라미터는 무엇인가요? 이 run에 연결된 메타데이터는 무엇인가요?

이 정보를 config 딕셔너리 (또는 유사한 객체)에 저장해 두고 필요할 때마다 꺼내 쓰는 방식은 매우 일반적인 워크플로입니다.

이 예제에서는 일부 하이퍼파라미터만 바뀌도록 두고 나머지는 직접 하드코딩합니다. 하지만 모델의 어떤 부분이든 config에 포함할 수 있습니다.

메타데이터도 함께 포함합니다. 여기서는 MNIST 데이터셋과 합성곱 아키텍처를 사용합니다. 나중에 예를 들어 같은 프로젝트에서 CIFAR에 완전연결 아키텍처를 사용하게 되면 이 정보가 run들을 구분하는 데 도움이 됩니다.

config = dict(
    epochs=5,
    classes=10,
    kernels=[16, 32],
    batch_size=128,
    learning_rate=0.005,
    dataset="MNIST",
    architecture="CNN")

이제 전체 파이프라인을 정의해 보겠습니다. 이는 모델 트레이닝에서 매우 일반적인 흐름입니다:

1. 먼저 모델과 관련 데이터, 그리고 옵티마이저를 make한 다음
2. 그에 맞게 모델을 train하고 마지막으로
3. 트레이닝이 어떻게 진행되었는지 확인하기 위해 test합니다.

아래에서 이 함수들을 구현하겠습니다.

def model_pipeline(hyperparameters):

    # wandb 시작
    with wandb.init(project="pytorch-demo", config=hyperparameters) as run:
        # run.config를 통해 모든 하이퍼파라미터에 액세스하여 로깅이 실행과 일치하도록 합니다.
        config = run.config

        # 모델, 데이터, 최적화 문제 생성
        model, train_loader, test_loader, criterion, optimizer = make(config)
        print(model)

        # 이를 사용하여 모델 트레이닝
        train(model, train_loader, criterion, optimizer, config)

        # 최종 성능 테스트
        test(model, test_loader)

    return model

여기서 표준 파이프라인과 다른 유일한 점은 이 모든 과정이 wandb.init()의 컨텍스트 내에서 이뤄진다는 것입니다. 이 함수를 호출하면 사용자 코드와 W&B 서버 사이에 통신 경로가 설정됩니다.

config 딕셔너리를 wandb.init()에 전달하면 그 정보가 모두 즉시 로깅되므로, 실험에 사용할 하이퍼파라미터 값을 어떻게 설정했는지 항상 알 수 있습니다.

선택해 로깅한 값이 실제로 모델에 사용되는 값과 항상 일치하도록 하려면, 객체의 run.config 사본을 사용하는 것을 권장합니다. 몇 가지 예시는 아래 make의 정의를 확인하세요.

참고: 저희는 코드를 별도의 프로세스에서 실행하므로, 저희 측에 문제가 생기더라도 (예를 들어 거대한 바다 괴물이 데이터 센터를 공격하더라도) 사용자 코드가 중단되지 않습니다. Kraken이 다시 깊은 바다로 돌아가듯 문제가 해결되면, wandb sync로 데이터를 로깅할 수 있습니다.

def make(config):
    # 데이터 생성
    train, test = get_data(train=True), get_data(train=False)
    train_loader = make_loader(train, batch_size=config.batch_size)
    test_loader = make_loader(test, batch_size=config.batch_size)

    # 모델 생성
    model = ConvNet(config.kernels, config.classes).to(device)

    # 손실 함수 및 옵티마이저 생성
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(
        model.parameters(), lr=config.learning_rate)

    return model, train_loader, test_loader, criterion, optimizer

​

데이터 로딩과 모델 정의

이제 데이터를 어떻게 로드할지와 모델이 어떤 형태인지 정의해야 합니다.

이 부분은 매우 중요하지만, wandb가 없을 때와 달라지는 점은 없으므로 여기서는 자세히 다루지 않겠습니다.

def get_data(slice=5, train=True):
    full_dataset = torchvision.datasets.MNIST(root=".",
                                              train=train,
                                              transform=transforms.ToTensor(),
                                              download=True)
    #  [::slice]로 슬라이싱한 것과 동일
    sub_dataset = torch.utils.data.Subset(
      full_dataset, indices=range(0, len(full_dataset), slice))

    return sub_dataset


def make_loader(dataset, batch_size):
    loader = torch.utils.data.DataLoader(dataset=dataset,
                                         batch_size=batch_size,
                                         shuffle=True,
                                         pin_memory=True, num_workers=2)
    return loader

보통 모델을 정의하는 부분이 가장 재미있습니다.

하지만 wandb를 쓴다고 해서 달라지는 건 없으니, 표준 ConvNet 아키텍처를 그대로 사용하겠습니다.

이 부분은 이것저것 바꿔 보면서 여러 실험을 해 보는 것도 주저하지 마세요 — 모든 결과는 wandb.ai에 로깅됩니다.

# 일반 신경망 및 합성곱 신경망

class ConvNet(nn.Module):
    def __init__(self, kernels, classes=10):
        super(ConvNet, self).__init__()

        self.layer1 = nn.Sequential(
            nn.Conv2d(1, kernels[0], kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, kernels[1], kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(7 * 7 * kernels[-1], classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

​

트레이닝 로직 정의하기

이제 model_pipeline에서 train을 어떻게 수행할지 정의할 차례입니다.

여기서는 wandb의 두 함수인 watch와 log를 사용합니다.

​

run.watch()로 그라디언트를 추적하고, 나머지는 모두 run.log()로 로깅하기

run.watch()는 모델의 그라디언트와 파라미터를 트레이닝 중 매 log_freq step마다 로깅합니다.

트레이닝을 시작하기 전에 이 함수를 호출하기만 하면 됩니다.

나머지 트레이닝 코드는 그대로입니다: 에포크와 배치를 순회하면서 forward 및 backward pass를 수행하고 optimizer를 적용합니다.

def train(model, loader, criterion, optimizer, config):
    # wandb에 모델의 그라디언트, 가중치 등을 추적하도록 지시합니다.
    run = wandb.init(project="pytorch-demo", config=config)
    run.watch(model, criterion, log="all", log_freq=10)

    # 트레이닝을 실행하고 wandb로 추적합니다
    total_batches = len(loader) * config.epochs
    example_ct = 0  # 지금까지 본 예제 수
    batch_ct = 0
    for epoch in tqdm(range(config.epochs)):
        for _, (images, labels) in enumerate(loader):

            loss = train_batch(images, labels, model, optimizer, criterion)
            example_ct +=  len(images)
            batch_ct += 1

            # 25번째 배치마다 메트릭 보고
            if ((batch_ct + 1) % 25) == 0:
                train_log(loss, example_ct, epoch)


def train_batch(images, labels, model, optimizer, criterion):
    images, labels = images.to(device), labels.to(device)

    # Forward pass ➡
    outputs = model(images)
    loss = criterion(outputs, labels)

    # Backward pass ⬅
    optimizer.zero_grad()
    loss.backward()

    # optimizer로 step 수행
    optimizer.step()

    return loss

유일한 차이는 로깅 코드입니다: 이전에는 메트릭을 터미널에 출력했다면, 이제는 같은 정보를 run.log()에 전달합니다.

run.log()는 키가 문자열인 딕셔너리를 받습니다. 이 문자열은 로깅할 객체를 식별하고, 해당 객체가 값이 됩니다. 현재 트레이닝의 몇 번째 step인지도 선택적으로 로깅할 수 있습니다.

참고: 저는 모델이 지금까지 본 예제 수를 사용하는 편입니다. 이렇게 하면 batch size가 달라도 비교하기가 더 쉽기 때문입니다. 물론 단순 step 수나 batch 개수를 사용해도 됩니다. 트레이닝 run이 더 길다면 epoch 기준으로 로깅하는 것도 의미가 있습니다.

def train_log(loss, example_ct, epoch):
    with wandb.init(project="pytorch-demo") as run:
        # loss와 에포크 번호를 로깅합니다
        # 여기서 메트릭을 W&B에 로깅합니다
        run.log({"epoch": epoch, "loss": loss}, step=example_ct)
        print(f"Loss after {str(example_ct).zfill(5)} examples: {loss:.3f}")

​

테스트 로직 정의

모델의 트레이닝이 끝나면 이를 테스트합니다: 예를 들어 프로덕션의 새로운 데이터에 실행해 보거나, 직접 엄선한 예제에 적용할 수 있습니다.

​

(선택 사항) run.save() 호출

지금은 모델의 아키텍처와 최종 파라미터를 디스크에 저장해 두기에도 좋은 시점입니다. 최대한의 호환성을 위해 모델을 Open Neural Network eXchange (ONNX) 형식으로 export하겠습니다.

해당 파일명을 run.save()에 전달하면 모델 파라미터가 W&B 서버에 저장되므로, 어떤 .h5 또는 .pb 파일이 어떤 트레이닝 run에 해당하는지 더 이상 놓치지 않게 됩니다.

모델을 저장하고, 버전 관리하고, 배포하는 더 고급 wandb 기능은 Artifacts 도구를 확인하세요.

def test(model, test_loader):
    model.eval()

    with wandb.init(project="pytorch-demo") as run:
        # 일부 테스트 예제에 모델 실행
        with torch.no_grad():
            correct, total = 0, 0
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

            print(f"Accuracy of the model on the {total} " +
                f"test images: {correct / total:%}")

            run.log({"test_accuracy": correct / total})

        # 교환 가능한 ONNX 형식으로 모델 저장
        torch.onnx.export(model, images, "model.onnx")
        run.save("model.onnx")

​

트레이닝을 실행하고 wandb.ai에서 메트릭을 실시간으로 확인하세요

이제 전체 파이프라인을 정의하고 몇 줄의 W&B 코드를 추가했으니, 완전히 추적되는 실험을 실행할 준비가 되었습니다.

몇 가지 링크를 안내해 드리겠습니다: 문서, 프로젝트의 모든 run을 정리해 보여주는 프로젝트 페이지, 그리고 이 run의 결과가 저장되는 Run 페이지입니다.

Run 페이지로 이동해 다음 탭을 확인하세요:

1. 차트, 트레이닝 전반에 걸쳐 모델의 그라디언트, 파라미터 값, loss가 로깅되는 곳입니다
2. 시스템, 디스크 I/O 사용량, CPU 및 GPU 메트릭(온도가 치솟는 것도 볼 수 있습니다) 등 다양한 시스템 메트릭이 포함되어 있습니다
3. 로그, 트레이닝 중 표준 출력으로 전송된 모든 내용의 복사본이 있습니다
4. 파일, 트레이닝이 완료되면 model.onnx를 클릭해 Netron model viewer에서 네트워크를 확인할 수 있습니다.

run이 완료되고 with wandb.init() 블록을 빠져나오면 셀 출력에도 결과 요약이 표시됩니다.

# 파이프라인으로 모델을 구축, 트레이닝 및 분석합니다
model = model_pipeline(config)

​

Sweeps로 하이퍼파라미터 테스트하기

이 예제에서는 하나의 하이퍼파라미터 세트만 살펴보았습니다. 하지만 대부분의 ML 워크플로에서 중요한 부분은 여러 하이퍼파라미터를 바꿔 가며 반복적으로 실험하는 것입니다.

W&B Sweeps를 사용하면 하이퍼파라미터 테스트를 자동화하고 가능한 모델과 최적화 전략의 공간을 탐색할 수 있습니다.

W&B Sweeps를 사용한 하이퍼파라미터 최적화를 보여주는 Colab 노트북을 확인해 보세요.

W&B에서 하이퍼파라미터 sweep을 실행하는 것은 매우 쉽습니다. 단 3단계만 거치면 됩니다:

1. sweep 정의: 검색할 파라미터, 검색 전략, 최적화 metric 등을 지정하는 딕셔너리 또는 YAML 파일을 만듭니다.

2. sweep 초기화: sweep_id = wandb.sweep(sweep_config)

3. sweep agent 실행: wandb.agent(sweep_id, function=train)

이것으로 하이퍼파라미터 sweep 실행이 끝납니다.

​

예제 갤러리

W&B로 추적하고 시각화한 프로젝트 예제를 Gallery →에서 확인해 보세요.

​

고급 설정

1. 환경 변수: 관리형 클러스터에서 트레이닝을 실행할 수 있도록 환경 변수에 API 키를 설정합니다.
2. 오프라인 모드: dryrun 모드를 사용해 오프라인으로 트레이닝한 뒤, 나중에 결과를 동기화합니다.
3. 온프레미스: 자체 인프라의 프라이빗 클라우드 또는 외부망과 격리된 서버에 W&B를 설치합니다. 학계 사용자부터 엔터프라이즈 팀까지 폭넓게 사용할 수 있는 로컬 설치 옵션을 제공합니다.
4. Sweeps: 경량 튜닝 도구로 하이퍼파라미터 검색을 빠르게 설정합니다.