[개발 환경] Wandb (Weights & Biases)

Wandb는 머신러닝 실험을 트래킹하고, 시각화해주는 MLOps 플랫폼이다.

1.

먼저, 터미널에서 아래 명령어를 실행해 패키지를 설치하고 로그인해야 한다.

pip install wandb
wandb login

위 명령어를 입력하면 API 키를 입력하라는 창이 뜰 것이다.
Wandb에 접속해서 회원가입을 하면, API key가 발급되는데, 이 key를 붙여넣기하면 된다.

2. 프로젝트 초기화

먼저, 학습 스크립트의 가장 첫 부분에서 wandb.init을 통해 wandb를 실행하고 하이퍼파라미터를 설정한다.

import wandb

wandb.init(
    project="my-first-dl-project", # 프로젝트 이름
    entity="my-team",              # 팀 워크스페이스 이름 (선택)
    name="experiment-1",           # 현재 실험의 이름 (run name)
    config={
        "learning_rate": 0.001,
        "epochs": 100,
        "batch_size": 32,
        "architecture": "ResNet"
    }
)
config = wandb.config

Wandb 대시보드에 접속하면 아래 그림과 같은 화면을 볼 수 있다.
코드에서 설정한 project 파라미터는 아래 그림의 왼쪽처럼 전체 실험을 묶어주는 프로젝트 이름이 되고, name 파라미터는 아래 그림의 오른쪽처럼 개별 실험 이름 (Run name)이 된다.

3. 데이터 logging

추적하고 싶은 지표는 학습 루프 안에서 wandb.log를 통해 딕셔너리 형태로 기록하면 된다.

for epoch in range(config.epochs):
    ...

    loss = 0.5
    acc = 85.0
    
    # 매 epoch마다 지표 기록
    wandb.log({
        "train_loss": loss,
        "train_acc": acc,
        "epoch": epoch
    })

이미지 logging

모델이 생성한 결과물, 시각화된 차트 등을 대시보드에서 바로 확인하고 싶을 때에는, 저장할 이미지를 wandb.Image() 객체로 감싼 다음 wandb.log() 안에 넣어주면 된다.

# 파일 경로를 이미지로 저장하는 법
wandb.log({"my_sample_image": wandb.Image("my_image.png", caption="원본 이미지")})

# NumPy 배열 형태를 이미지로 저장하는 법
random_image_array = np.random.randint(255, size=(100, 100, 3), dtype=np.uint8)
wandb.log({"generated_noise": wandb.Image(random_image_array, caption="Numpy로 만든 노이즈")})

# PIL(Pillow) 이미지 객체를 이미지로 저장하는 법
pil_img = Image.open("my_image.png")
wandb.log({"pil_image_example": wandb.Image(pil_img, caption="PIL 이미지")})

# 여러 이미지를 리스트에 담을 수도 있음
image_list = [
    wandb.Image("image1.png", caption="이미지 1"),
    wandb.Image("image2.png", caption="이미지 2"),
    wandb.Image("image3.png", caption="이미지 3")
]
wandb.log({"image_gallery": image_list})

4. 모델 및 결과물 저장 (wandb.Artifact)

학습이 끝난 모델의 가중치나 데이터셋을 wandb.Artifact를 통해 버전별로 저장할 수 있다.

# wandb 서버에 업로드할 빈 상자 생성 (데이터셋이면 type='dataset')
artifact = wandb.Artifact('model-weights', type='model')

# 빈 상자에 파일을 담음
artifact.add_file('best_model.pth')

# wandb 대시보드로 전송
wandb.log_artifact(artifact)

5. 실험 종료

단일 스크립트 내에서 여러 번의 실험을 돌릴 때는 반드시 wandb.finish를 통해 종료 선언을 해야 한다.

# wandb 세션 종료
wandb.finish()