PyTorch

기계학습 실험 트래킹, 데이터셋 버전 관리 및 프로젝트 협업을 위해 W&B를 사용해 보세요.

​

이 노트북에서 다루는 내용

PyTorch 코드와 W&B를 통합하여 파이프라인에 실험 트래킹 기능을 추가하는 방법을 보여드립니다.

# 라이브러리 임포트
import wandb

# 새로운 실험 시작
with wandb.init(project="new-sota-model") as run:
 
    # config를 사용하여 하이퍼파라미터 사전을 캡처
    run.config = {"learning_rate": 0.001, "epochs": 100, "batch_size": 128}

    # 모델 및 데이터 설정
    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의 완전 연결(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를 통해 모든 HP에 액세스하여 로그 기록과 실행이 일치하도록 함
        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 정의에서 몇 가지 예시를 확인하세요.

참고: 당사는 코드를 별도의 프로세스에서 실행하도록 관리하여, 당사 측의 문제(예: 거대 바다 괴물이 데이터 센터를 공격하는 경우 등)가 사용자의 코드를 중단시키지 않도록 합니다. 문제가 해결되면(예: 크라켄이 심해로 돌아간 후), 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 방법을 지정할 차례입니다.

여기서는 watch와 log, 두 가지 wandb 함수가 사용됩니다.

​

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)
    
    # Forward 패스 ➡
    outputs = model(images)
    loss = criterion(outputs, labels)
    
    # Backward 패스 ⬅
    optimizer.zero_grad()
    loss.backward()

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

    return loss

유일한 차이점은 로그 기록 코드에 있습니다: 이전에는 메트릭을 터미널에 출력하여 보고했다면, 이제는 동일한 정보를 run.log()에 전달합니다.

run.log()는 문자열을 키로 하는 사전을 인자로 받습니다. 이 문자열은 로그되는 오브젝트를 식별하며, 그 값이 로그 데이터가 됩니다. 선택적으로 현재 트래닝 중인 step을 함께 로그할 수도 있습니다.

참고: 저는 모델이 학습한 샘플 수를 사용하는 것을 선호합니다. 이렇게 하면 배치 크기가 달라도 비교하기가 더 쉽기 때문입니다. 하지만 원시 스텝이나 배치 수를 사용할 수도 있습니다. 트래닝 기간이 긴 경우 epoch 단위로 로그를 기록하는 것도 합리적입니다.

def train_log(loss, example_ct, epoch):
    with wandb.init(project="pytorch-demo") as run:
        # 손실 및 에포크 번호 로그 기록
        # 여기서 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을 정리해 보여주는 Projects 페이지, 그리고 이 run의 결과가 저장될 Run 페이지입니다.

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

1. Charts: 트래닝 전반에 걸쳐 모델의 그레이디언트, 파라미터 값, 손실이 로그되는 곳입니다.
2. System: 디스크 I/O 사용량, CPU 및 GPU 메트릭(온도 상승 확인) 등 다양한 시스템 메트릭을 포함합니다.
3. Logs: 트래닝 중 표준 출력(standard out)으로 전송된 모든 내용의 복사본이 있습니다.
4. Files: 트래닝이 완료된 후 model.onnx를 클릭하여 Netron 모델 뷰어로 네트워크를 시각화할 수 있습니다.

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)

하이퍼파라미터 스윕을 실행하는 데 필요한 과정은 이것이 전부입니다.

​

예제 갤러리

W&B로 트래킹되고 시각화된 프로젝트 예제를 Gallery →에서 살펴보세요.

​

고급 설정

1. 환경 변수: 관리형 클러스터에서 트래닝을 실행할 수 있도록 환경 변수에 API 키를 설정합니다.
2. 오프라인 모드: dryrun 모드를 사용하여 오프라인에서 트래닝하고 나중에 결과를 동기화합니다.
3. 온프레미스: 자체 인프라의 프라이빗 클라우드나 에어갭 서버에 W&B를 설치합니다. 학계부터 엔터프라이즈 팀까지 모두를 위한 로컬 설치 옵션을 제공합니다.
4. Sweeps: 튜닝을 위한 가벼운 툴을 사용하여 하이퍼파라미터 검색을 신속하게 설정합니다.