PyTorch

Colab에서 실행해 보기

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

​

이 노트북에서 다루는 내용

이 노트북에서는 PyTorch 코드에 W&B를 인테그레이션해 파이프라인에 실험 추적 기능을 추가하는 방법을 보여줍니다.

# 라이브러리를 import합니다
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")]

​

0단계: W&B 설치

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

!pip install wandb onnx -Uq

​

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

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

W&B를 처음 사용하는 경우, 표시되는 링크에서 무료 계정을 먼저 생성해야 합니다.

import wandb

wandb.login()

​

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

​

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

프로그램 상에서 우리가 가장 먼저 하는 일은 실험을 정의하는 것입니다: 하이퍼파라미터는 무엇인가요? 이 run에 연결된 메타데이터는 무엇인가요?

이 정보를 config 사전(또는 비슷한 객체)에 저장해 두었다가 필요할 때 참조하는 것은 꽤 일반적인 워크플로우입니다.

이 예제에서는 소수의 하이퍼파라미터만 변경 가능하게 두고 나머지는 코드에 직접 지정합니다. 하지만 모델의 어떤 부분이든 config의 일부가 될 수 있습니다.

또한 일부 메타데이터도 포함합니다: MNIST 데이터셋과 합성곱 아키텍처를 사용하고 있습니다. 나중에 같은 프로젝트에서 예를 들어 CIFAR에 완전 연결(fully connected) 아키텍처를 사용해 작업하더라도, 이 정보가 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의 컨텍스트 안에서 수행된다는 점입니다. 이 함수를 호출하면 사용자의 코드와 우리의 서버 사이에 통신 채널이 설정됩니다.

config 사전을 wandb.init에 넘기면 그 정보가 즉시 우리에게 로깅되므로, 실험에서 어떤 하이퍼파라미터 값을 사용하도록 설정했는지 항상 알 수 있습니다.

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

추가 참고: 우리는 코드가 별도의 프로세스에서 실행되도록 주의 깊게 설계했기 때문에, 우리 쪽에서 발생하는 어떤 문제 (예를 들어 거대한 바다 괴물이 우리 데이터 센터를 공격하는 경우)로 인해 사용자의 코드가 크래시 나는 일은 없습니다. 문제가 해결된 후, 예를 들어 크라켄이 심해로 돌아가면, 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 스텝마다 모델의 그래디언트와 파라미터를 기록합니다.

트레이닝을 시작하기 전에 이것만 호출해 주면 됩니다.

나머지 트레이닝 코드는 그대로입니다: 에포크와 배치를 반복하면서 forward와 backward 패스를 실행하고 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)
    
    # 순전파 ➡
    outputs = model(images)
    loss = criterion(outputs, labels)
    
    # 역전파 ⬅
    optimizer.zero_grad()
    loss.backward()

    # 옵티마이저 스텝
    optimizer.step()

    return loss

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

run.log()는 문자열을 key로 사용하는 사전을 인자로 받습니다. 이 문자열들은 로깅되는 객체를 식별하고, 해당 객체들이 value를 이룹니다. 원한다면 현재 트레이닝에서 몇 번째 step인지도 함께 로깅할 수 있습니다.

추가 참고: 저는 모델이 지금까지 본 예제 개수를 사용하는 방법을 선호합니다. 이렇게 하면 배치 크기가 달라져도 비교하기가 더 쉽기 때문입니다. 하지만 원한다면 step 값이나 배치 개수를 그대로 사용할 수도 있습니다. 트레이닝 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) format으로 export하겠습니다.

그 파일 이름을 run.save()에 전달하면 모델 파라미터가 W&B의 서버에 저장됩니다. 이제 어떤 .h5나 .pb 파일이 어떤 트레이닝 run에 해당하는지 헷갈릴 일이 없습니다.

모델을 저장, 버전 관리, 배포하기 위한 더 고급 wandb 기능은 Artifacts tools를 참조하세요.

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 코드를 살짝 추가했으니, 완전히 추적되는 실험을 실행할 준비가 되었습니다.

다음과 같은 몇 가지 링크를 제공해 드립니다: W&B 문서, 프로젝트의 모든 runs를 정리해 주는 Project 페이지, 그리고 이 run의 결과가 저장될 Run 페이지입니다.

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

1. Charts 탭: 트레이닝 전체에 걸쳐 모델 그래디언트, 파라미터 값, 그리고 loss가 로깅됩니다.
2. System 탭: Disk I/O 사용량, CPU 및 GPU 메트릭(온도가 치솟는 것도 볼 수 있습니다) 등 다양한 시스템 메트릭이 포함됩니다.
3. Logs 탭: 트레이닝 중 표준 출력으로 보내진 모든 내용의 복사본이 저장됩니다.
4. Files 탭: 트레이닝이 완료되면 model.onnx를 클릭해 Netron model viewer에서 네트워크를 확인할 수 있습니다.

run이 완료되고 with wandb.init 블록이 종료되면, 셀 출력에 결과 요약도 함께 표시됩니다.

# 파이프라인으로 모델 빌드, 학습 및 분석
model = model_pipeline(config)

​

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

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

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

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

W&B로 하이퍼파라미터 스윕을 실행하는 것은 매우 쉽습니다. 다음 3가지 간단한 단계만 거치면 됩니다.

1. 스윕 정의하기: 검색할 파라미터, 검색 전략, 최적화 메트릭 등을 지정하는 사전 또는 YAML 파일을 만들어 정의합니다.

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

3. 스윕 에이전트 실행하기:
   wandb.agent(sweep_id, function=train)

이렇게 하면 하이퍼파라미터 스윕 실행에 필요한 모든 단계가 완료됩니다.

​

예제 갤러리

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

​

고급 설정

1. Environment variables: 관리형 클러스터에서 트레이닝을 실행할 수 있도록 환경 변수에 API 키를 설정하세요.
2. Offline mode: dryrun 모드를 사용해 오프라인으로 트레이닝하고, 나중에 결과를 동기화하세요.
3. On-prem: 자체 인프라의 프라이빗 클라우드 또는 에어갭 환경의 서버에 W&B를 설치하세요. 학계부터 엔터프라이즈 팀까지 모두를 위한 로컬 설치 옵션을 제공합니다.
4. Sweeps: 튜닝을 위한 가벼운 도구를 사용해 하이퍼파라미터 탐색을 빠르게 설정하세요.