[WandB] Weights & Biases Tutorial (영상 처리)

es.Seong·2024년 3월 12일
0
post-thumbnail

WandB 주소
https://kr.wandb.ai/

WandB(Weights & Biases)는 ML 실험에 사용한 하이퍼파라미터나 시각화자료, 예측 결과 등을 추적 및 관리하는 툴이다.
연구를 진행하며 무수히 많은 실험을 할텐데 이를 모두 수기로 기록하기에는 오기입, 누락 등이 발생하여 딥러닝의 재현성(Reproducibility)에 문제가 발생할 수 있다. 제안 모형을 튜닝하는 과정에서 최적의 파라미터를 잊어버린다면 큰 시간낭비로 이어진다.

WandB (대부분 완디비라고 말하는 것 같다.)를 활용하며 이미지 처리에 어떻게 사용하는지 공유하려 합니다.

WandB 회원 가입 과정은 생략하였습니다.

개발 환경
OS : MAC OS
Python : 3.9.13
Pytorch : 2.2.0

1. Login

WandB 설치
pip install wandb

pip 명령어를 통해 wandb 라이브러리를 설치했다면, 터미널에 wandb login을 입력해보자.
다음과 같은 멘트가 출력될 것이다.

wandb login

wandb: You can find your API key in your browser here: https://app.wandb.ai/authorize wandb: Paste an API key from your profile and hit enter:

해당 문구가 출력된다면 https://app.wandb.ai/authorize 에 접속하여 API Key 값을 복사해주자.
그리고 다시 터미널로 돌아와 Ctrl+V 해주면 인증이 완료된다.

붙여넣기를 해도 키값은 보이지 않는다. 붙여넣고 바로 엔터를 눌러주면 정상적으로 접속된다.

이제 터미널에서 해야할 작업은 없고 코드에서 wandb 라이브러리를 import하여 진행하면된다.

import wandb

# wandb 설정
wandb.login()
wandb.init(project="Project_Name",name="ModelName",tags=["tag1","tag2","tag3"])

wandb.login() : .py, .ipynb 파일내에서 로그인을 진행한다.
wandb.init() : wandb 초기설정을 진행하는 메소드이다.
메소드 내 파라미터는 주로 project, name, tags를 주로 사용한다.
project : 연구 프로젝트명을 입력한다.
name : 연구 내 사용하는 모형 이름이나 모형 특징을 입력한다.
tags : 필터링에 사용할 문자열.

2. Hyperparameter 정의

wandb에서는 모형에 사용할 하이퍼파라미터를 저장하는 기능을 제공한다.
아래 코드를 한 번 살펴보자.

args = {
    "learning_rate": 0.0002,
    "epochs": 3,
    "batch_size": 32,
    "device": DEVICE,
    "image_size": 128,
    "channels": 1,
    "beta1" : 0.5,
    "beta2" : 0.999,
    "noise" : "gauss"
    
}
wandb.config.update(args)

args라는 딕셔너리의 key, value에 저장할 하이퍼파라미터 이름과 사용할 값을 순서대로 입력한다.
주로 학습률, 반복 수, 배치 사이즈, 입력 이미지 크기 등을 저장하면 된다.

wandb.config.update(args) 코드를 통해 wandb에 해당 하이퍼파라미터들을 모두 기록할 수 있다.

3. 데이터 저장

3-1. 이미지 저장

아래 코드는 파이토치로 정의한 Train Dataset에서 이미지 다섯장을 wandb에 저장하는 코드이다.

wandb_list = []
for i in range(5):
    plt.subplot(1, 5, i+1)
    plt.imshow(train_dataset[i].squeeze(), cmap='gray')
    plt.title(f'Index: {i}')
    plt.axis('off')
    wandb_list.append(wandb.Image(train_dataset[i].squeeze(), f'Index: {i}'))
wandb.log({"Train Dataset Sample": wandb_list})
plt.show()

wandb_list변수에 이미지 다섯장 픽셀 값을 모두 저장한다.
wandb.log 메소드는 정의한 변수를 wandb 내 Dashboard에 저장 및 출력하는 변수다.

아래 이미지는 WandB 내 출력된 결과이다.

Media 부분에 다섯장의 이미지가 저장된 것을 확인할 수 있다.

3-2. 학습 중 Metric 저장

# Epoch Step 별 데이터 저장
#학습 코드 생략

wandb.log({"epoch":epoch,"Loss_D":errD.item(),"Loss_G":g_loss.item(),"D(x)":D_x,"D(G(z))":D_G_z1,"D(G(z))":D_G_z2,"psnr":psnr})

해당 예시는 GAN 모형을 학습 할 때 출력되는 값들이다. Epoch 수 만큼 학습이 진행되면서 값들이 어떻게 변하는지 확인하기 위해 wandb.log를 사용하여 Dashboard에 출력할 수 있다.


wandb.log 내 선언된 모든 변수가 각각의 plot에 저장된 것을 확인할 수 있다.

3-3. 영상 저장

이미지뿐만 아니라 .gif와 같은 영상 또한 저장이 가능하다. 단, gif 변수를 바로 저장하지는 못하고, 로컬에 파일을 저장한 후 wandb에 저장해야한다.

import matplotlib.animation as animation
from IPython.display import HTML
fig = plt.figure(figsize=(15,15))
plt.axis("off")
ims = [[plt.imshow(np.transpose(i,(1,2,0)), animated=True)] for i in img_list]
ani = animation.ArtistAnimation(fig, ims, interval=1000, repeat_delay=1000, blit=True)

ani.save('./test.gif', writer='imagemagick', fps=60)
wandb.log({"test":wandb.Video('./test.gif',fps=60,format="gif")})

HTML(ani.to_jshtml())

이미지 데이터를 이어붙여서 영상으로 만드는 코드이다.
gif가 저장된 코드는 ani 변수이며 save 메소드를 통해 로컬에 저장한다.
저장한 파일을 wandb.log({"test":wandb.Video('./test.gif',fps=60,format="gif")})를 사용하여 wandb에 저장한다. 동영상 파일은 wandb.Image가 아닌 wandb.Video를 사용해야한다.

3-4. 모델 가중치 저장

모델 가중치 저장은 영상데이터 저장 방법과 동일하다. 단, wandb.log를 사용하지 않고 wandb.save(weight)만 입력하면된다.

저장된 가중치를 WandB 페이지에서 확인하는 방법은 다음과 같다.
좌측 Files 클릭

정의한 경로로 이동
저장 경로는 코드 내 정의한 경로와 동일하게 생성된다.

4. Finish

모든 과정이 끝이 났다면 finish 메소드를 통해 종료하면 된다. wandb와의 연동이 종료된다.

#wandb 접속 종료 및 저장
wandb.finish()

글에서 다루진 않았지만 WandB Dashboard에는 GPU, CPU 할당량, 사용량 같은 시스템적인 부분도 모니터링도 지원하여 명령어를 통해 확인해야하는 수고로움을 덜어준다.
또한 init에서 name 변수로 선언한 케이스 별로 Metic 비교도 되어 연구나 해커톤 등에 아주 유용하게 사용할 것 같다.

profile
Graduate student at Pusan National University, majoring in Artificial Intelligence

0개의 댓글