WandB 사용법

지난 경진대회에서 시간이 없어서 시도해보지 못한 sweep 기능을 적용하면서 한번 정리해보려고 한다.

1. WandB Sweep이란?

WandB의 Sweep 기능은 하이퍼파라미터 최적화 실험을 자동화하는 도구다. learning_rate, batch_size, dropout 등 다양한 하이퍼파라미터 조합을 자동으로 탐색하여 최적의 조합을 찾아주는 기능을 제공한다.

Sweep의 장점

• 자동화된 하이퍼파라미터 튜닝: 일일이 실험을 돌릴 필요 없이 설정된 조합에 따라 실험이 자동으로 진행된다.
• 효율적인 탐색: 다양한 하이퍼파라미터 조합을 효과적으로 탐색하고 최적의 성능을 찾아낼 수 있다.
• 결과 시각화: Parallel 플롯과 같은 시각화 도구로 하이퍼파라미터와 성능 간 관계를 직관적으로 파악할 수 있다.

 

2. WandB Sweep의 옵션

WandB Sweep을 설정할 때 사용할 수 있는 주요 옵션은 다음과 같다.

Sweep Configuration Options

method

• 하이퍼파라미터 조합을 탐색하는 방법을 설정한다.
• grid: 모든 가능한 조합을 시도하는 방식이다. 가능한 모든 조합을 시도하므로 비용이 많이 들 수 있다.
• random: 랜덤하게 조합을 선택하여 탐색하는 방식으로 빠르게 최적의 조합을 찾을 수 있다.
• bayes: 베이지안 최적화 방식을 사용하여 효율적으로 최적의 조합을 찾는다.

metric

최적화할 목표 지표를 설정한다. 예를 들어, "loss"를 최소화하거나 "accuracy"를 최대화하는 방식으로 설정할 수 있다.

metric: 
  name: "valid_f1"  # 최적화할 목표 지표
  goal: "maximize"  # maximize 또는 minimize로 설정

 

parameters

하이퍼파라미터 값의 범위를 설정한다. 각 하이퍼파라미터에 대해 values 리스트로 시도할 값을 나열하거나 min과 max 범위를 지정할 수 있다.

parameters:
  learning_rate:
    values: [0.001, 0.0005, 0.0001]
  batch_size:
    min: 8
    max: 64
  epochs:
    values: [5, 10, 15]

 

 

3. WandB Sweep 코드 적용

이제 실제로 WandB Sweep을 사용하여 하이퍼파라미터 최적화 실험을 자동화해보자.

아래 코드는 실제 경진대회 코드에 WandB Sweep을 설정하고 학습 코드를 정의하여 다양한 하이퍼파라미터 조합으로 모델을 학습하는 과정이다.

Step 1: Sweep Configuration 정의하기

Sweep 설정을 정의한다. 이 설정은 method, metric, parameters의 범위를 포함한다.

import wandb

sweep_config = {
    "method": "bayes",  # 또는 "grid", "random"
    "metric": {
        "name": "valid_loss",
        "goal": "minimize"  # val_loss 기준으로 최소화하도록 설정
    },
    "parameters": {
        "learning_rate": {
            "values": [1e-3, 1e-4, 5e-4]
        },
        "batch_size": {
            "values": [8, 16, 32]
        },
        "epochs": {
            "values": [10, 20]
        },
        # 추가로 다른 하이퍼파라미터를 설정할 수 있다.
    }
}

sweep_id = wandb.sweep(sweep_config, project="sweep_efficientnet_b0")

Step 2: 모델 학습 함수 정의하기

이제 실제 모델 학습하는 코드에 wandb sweep을 녹여보자.

model_name = 'tf_efficientnet_b0'

def train_with_sweep():
    wandb.init()  # wandb.init()을 먼저 호출하여 설정 초기화
    config = wandb.config  # wandb.config를 통해 하이퍼파라미터 불러오기
    k_folds = 3
    skf = StratifiedKFold(n_splits=k_folds, shuffle=True, random_state=2024)

    for fold, (train_idx, val_idx) in enumerate(skf.split(trn_dataset, trn_dataset.targets)):
        print(f'Fold {fold + 1}/{k_folds}')

        train_subset = Subset(trn_dataset, train_idx)
        val_subset = Subset(trn_dataset, val_idx)
        train_loader = DataLoader(train_subset, batch_size=config.batch_size, shuffle=True, num_workers=num_workers)
        val_loader = DataLoader(val_subset, batch_size=config.batch_size, shuffle=False, num_workers=num_workers)

        model = timm.create_model(model_name, pretrained=True, num_classes=17).to(device)
        optimizer = Adam(model.parameters(), lr=config.learning_rate)
        loss_fn = nn.CrossEntropyLoss()
        scheduler = CosineAnnealingLR(optimizer, T_max=10, eta_min=1e-4)

        for epoch in range(config.epochs):
            train_metrics = train_one_epoch(train_loader, model, optimizer, loss_fn, device)
            val_metrics = validate_one_epoch(val_loader, model, loss_fn, device)
            scheduler.step()

            wandb.log({
                "train_loss": train_metrics['train_loss'],
                "train_acc": train_metrics['train_acc'],
                "train_f1": train_metrics['train_f1'],
                "valid_loss": val_metrics['val_loss'],
                "valid_acc": val_metrics['val_acc'],
                "valid_f1": val_metrics['val_f1'],
                "epoch": epoch
            })

        wandb.finish()

여기서 기존에는 model_name을 'efficientnet_b0'으로 작성했는데 Unknown model 에러가 발생했다.

Run hix0axqq errored:
Traceback (most recent call last):
    File "/opt/conda/lib/python3.10/site-packages/wandb/agents/pyagent.py", line 306, in _run_job         self._function()
    File "/tmp/ipykernel_2926515/2117931372.py", line 98, in train_with_sweep
        model = timm.create_model(model_name, pretrained=True, num_classes=17).to(device)
    File "/opt/conda/lib/python3.10/site-packages/timm/models/_factory.py", line 113, in create_model
        raise RuntimeError('Unknown model (%s)' % model_name)
RuntimeError: Unknown model (efficientnet-b0)

 

이 오류는 timm.create_model 함수가 efficientnet-b0 모델을 인식하지 못할 때 발생한다고 한다. timm 라이브러리에서 지원하는 모델 이름이 정확해야 하는데 이 경우 정확한 모델 이름은 tf_efficientnet_b0라고 한다. efficientnet-b0는 일반적인 명명 방식이지만 timm에서는 tf_efficientnet_b0와 같은 이름을 사용한다고 해서 이름을 바꿔주고 실행하니 에러가 해결되었다.

Step 3: Sweep 실행하기

이제 wandb.agent를 사용하여 Sweep을 실행해주자.

wandb.agent(sweep_id, function=train_with_sweep)

이 코드를 실행하면 WandB가 자동으로 sweep_config에 설정한 하이퍼파라미터 조합을 사용해 train_with_sweep 함수를 반복적으로 실행한다. 이 과정에서 wandb가 지정한 하이퍼파라미터 조합에 따라 모델을 학습하고 성능을 평가하고 각 실험의 결과가 WandB 대시보드에 기록된다.

 

wandb.agent의 동작 방식

1. 하이퍼파라미터 조합 생성: sweep_config에서 설정한 하이퍼파라미터 공간 내에서 조합을 생성한다. 지금 같은 상황에서는 learning_rate, batch_size, epochs의 다양한 조합을 선택한다.
2. 학습 함수 실행: 각 조합에 대해 train_with_sweep 함수를 실행하여 모델을 학습한다. wandb.config를 통해 train_with_sweep 함수 내부에서 해당 조합의 하이퍼파라미터를 자동으로 가져와 사용할 수 있다.
3. 실험 결과 기록: 각 조합에 대한 학습과 평가가 완료되면 wandb.log()로 로그를 남기고 WandB 대시보드에서 실험의 성능을 비교할 수 있다.
4. 최적의 조합 선택: 설정된 기준(지금은 valid_f1을 maximize)에 따라 최적의 하이퍼파라미터 조합을 자동으로 탐색한다.

따라서 wandb.agent(sweep_id, function=train_with_sweep)을 실행하면 여러 조합에 대해 자동으로 모델이 돌아가고 별도로 추가적인 실행이나 수동 설정 없이도 wandb가 최적의 조합을 찾도록 도와준다.

 

위 코드를 실행하면 이렇게 wandb가 정상작동하는 로그를 볼 수 있다.

sweep 실행 성공 모습

 

4. WandB Sweep 결과 분석

모델 학습이 완료되면 WandB의 Sweep 대시보드에서 결과를 분석할 수 있다.

 

Parallel Coordinates Plot

1. Charts 메뉴에서 Parallel Coordinates Plot을 선택한다. (근데 아마 기본으로 노출될거다.)
2. 각 축에 learning_rate, batch_size, epochs, valid_f1와 같은 하이퍼파라미터와 성능 지표를 추가해 시각화한다.
3. 최적의 성능을 나타내는 조합을 확인하고 하이퍼파라미터와 성능 간의 관계를 분석할 수 있다.
4. 이렇게 실행하면 다양한 시각화 툴을 제공해준다.
   
   

WandB sweep 시각화 결과

 

 

그런데 사실 코드를 진행하면서 너무오래걸리길 이것저것 찾아보니 sweep을 사용할때 k 교차검증을 사용하면 시간이 엄청 오래걸린다는것을 발견했다.

 

train_with_sweep() 함수가 매번 k_folds = 3에 따라 3개의 폴드로 교차 검증을 수행하고 매번 모든 하이퍼파라미터 조합에 대해 실행된다. 이렇게 하면 하이퍼파라미터 조합이 많아질수록 전체 실행 횟수가 기하급수적으로 증가해서 시간이 많이 걸릴 수 있다.

 

Sweep과 교차 검증을 함께 사용하는 경우의 문제점

일반적으로 WandB Sweep은 최적의 하이퍼파라미터 조합을 찾기 위해 다양한 조합을 테스트하는데 교차 검증까지 포함할 경우 계산 비용이 매우 커진다. 그래서 Sweep과 교차 검증을 함께 사용하는 것은 큰 규모의 실험에서는 비효율적일 수 있다.

대신, Sweep을 사용할 때 교차 검증을 생략하거나 특정 하이퍼파라미터 조합에서 교차 검증을 한 번만 수행하는 방식을 고려할 수도 있다.

 

예를 들어, k_folds=1로 설정해서 한 번의 학습과 검증만 수행하거나, k_folds를 다른 코드에서 설정해 최적 조합을 찾은 후 교차 검증을 수행하는 방식을 사용할 수도 있다.

교차 검증 없이 한 번의 학습과 검증만 수행하도록 train_with_sweep() 함수를 단순하게 바꿔봤다.

def train_with_sweep():
    wandb.init()
    config = wandb.config

    # 단일 학습 및 검증 데이터셋 분할 (교차 검증 제거)
    train_size = int(0.8 * len(trn_dataset))
    val_size = len(trn_dataset) - train_size
    train_subset, val_subset = random_split(trn_dataset, [train_size, val_size])

    train_loader = DataLoader(train_subset, batch_size=config.batch_size, shuffle=True, num_workers=num_workers)
    val_loader = DataLoader(val_subset, batch_size=config.batch_size, shuffle=False, num_workers=num_workers)

    # 모델 및 옵티마이저 설정
    model = timm.create_model(model_name, pretrained=True, num_classes=17).to(device)
    optimizer = Adam(model.parameters(), lr=config.learning_rate)
    loss_fn = nn.CrossEntropyLoss()
    scheduler = CosineAnnealingLR(optimizer, T_max=10, eta_min=1e-4)

    for epoch in range(config.epochs):
        train_metrics = train_one_epoch(train_loader, model, optimizer, loss_fn, device)
        val_metrics = validate_one_epoch(val_loader, model, loss_fn, device)
        scheduler.step()

        wandb.log({
            "train_loss": train_metrics['train_loss'],
            "train_acc": train_metrics['train_acc'],
            "train_f1": train_metrics['train_f1'],
            "valid_loss": val_metrics['val_loss'],
            "valid_acc": val_metrics['val_acc'],
            "valid_f1": val_metrics['val_f1'],
            "epoch": epoch
        })

    wandb.finish()

 

그게 아니면 Sweep을 먼저 실행해서 최적의 하이퍼파라미터 조합을 찾을수도 있다. 그때 교차 검증을 생략하고 단일 훈련/검증 분할로 간단하게 모델을 학습한다. 그리고 최적의 하이퍼파라미터 조합을 찾은 후에 그 조합을 사용해 별도로 교차 검증을 수행해서 모델의 일반화 성능을 평가할수도 있다.