W&B와 PyTorch 통합
W&B와 PyTorch 통합
PyTorch로 모델을 학습할 때 W&B를 붙이면, 하이퍼파라미터와 메트릭을 한곳에서 추적할 수 있어요. wandb.init으로 실험을 정의하고, 학습 루프 안에서 run.log()로 측정값을 남기는 패턴이 핵심이에요.
출처: W&B PyTorch 튜토리얼
PyTorch 통합의 기본 예시를 보면, config 딕셔너리로 하이퍼파라미터를 캡처한 뒤 wandb.init으로 새 실험을 시작해요.
import wandb
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") # 모델 저장
run.watch(model)을 호출하면 모델의 그라디언트와 파라미터까지 추적할 수 있어요. run.log()는 학습 루프 안에서 매 배치의 메트릭을 기록하고, run.save()로 학습이 끝난 모델 파일을 대시보드에 올려둘 수 있죠.
MNIST 예시에서는 device 설정부터 파이프라인이 잘 정리돼 있어요.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
하이퍼파라미터를 config로 넘기면, 학습 실행이 실제로 쓴 값과 로깅이 항상 일치해요. W&B의 config는 '실행 메타데이터' 역할을 해서, 나중에 여러 실험을 비교할 때 어떤 설정이었는지 바로 알 수 있게 해 주죠. PyTorch 외에도 Hugging Face 등 다양한 라이브러리 통합이 준비돼 있어요.