PyTorch

# 라이브러리를 임포트하세요
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")

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

# 결정론적 동작 보장
seed = 42
torch.backends.cudnn.deterministic = True
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# 디바이스 설정
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")]

!pip install wandb onnx -Uq

import wandb

wandb.login()

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

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

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

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

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

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

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

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}")

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")

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