[파이토치] 6. TensorBoard와 WandB

서론

 

스타트업에서는 다양한 가설과 이를 증명할 수 있는 기능을 만들어 실험을 진행하고 분석한 결과를 바탕으로 기능을 개선하는 형태로 제품을 만들어갑니다. 이 때의 실험에 대한 결과 분석이 매우 중요해서 다양한 로그와 지표들을 설정하고 실험 과정에서 이를 트래킹합니다. 딥러닝 모델을 만드는 것 역시도 제품을 개선하는 것과 유사하게 수많은 이터레이션과 그 이터레이션에서 레슨을 얻는 것이 필요한 작업입니다. 

 

어떤 것들을 실험 과정에서 기록해야할까요?

 

1. Loss & Accuracy가 에폭마다 잘 떨어지는가?

2. weight, bias가 Descent Vanishing으로 학습이 안되지는 않는가?

3. 우리가 해결하려는 문제의 차원들을 모델이 임베딩을 통해서 적절하게 구분하였는가?

4. 실험 중간중간 결과물들을 조회할 수 있는가?

 

하지만 그동안 패운 파이토치 관련 코드에서는 관련 코드가 없었는데 오늘은 이러한 기능을 제공해주는 Tensorboard와 WandB에 대해서 소개해보려고 한다.

 

1. TensorBoard

텐서보드는 TensorFlow의 공식 시각화 툴킷이지만, 파이토치에서도 사용될 수 있도록 확장되었다. 기본적으로 학습 과정에서의 주요 지표 변화를 트래킹하거나 학습 과정을 이해하는데 필요한 여러 시각화 기능들을 제공한다.

 

 

1. tensorboard 불러오기

# 설치
!pip install tensorboard==2.12.3 -q


#로그 파일 관리
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter("./runs/tutorial")

 

*tensorboard에서 저장된 로그 파일을 반영하는 데에 업로드하는 주기가 존재해서, 바로 반영이 안되면 잠깐 기다리고 다시 셀을 실행해주면 반영이 된다.

 

이제부터는 코랩에서 텐서보드 결과를 보기 위해서는 아래 코드를 실행해야한다.

#처음 로드할 때의 코드
# IPython 환경에서 tensorboard 확장 기능을 로드하는 역할을 합니다.
%load_ext tensorboard

# ./runs/tutorial 위치에 저장된 로그를 위치로 tensorboard를 실행합니다.
%tensorboard --logdir ./runs/tutorial


#이후 로드할 때의 코드
%reload_ext tensorboard
%tensorboard --logdir ./runs/tutorial

 

 

2. add_graph : 모델 구조 시각화하기

model = CNN(num_classes = 10, dropout_ratio = 0.2)
model.weight_initialization()

for images, labels in train_dataloader:
  break

# 모델의 구조를 로깅합니다. 'GRAPHS' 탭을 클릭하여 확인할 수 있습니다.
# add_graph에서 input_to_model을 인자로 넣어, 모델의 거치면서 tensor의 차원 변화를 알 수 있습니다.
writer.add_graph(model = model, input_to_model = images)

 

 

우리가 모델을 설계하는 과정에서 모델이 복잡해질수록 이 구조를 이해하는 것이 매우 어려워진다. tensorboard에서는 add_graph를 통해서 model과 input_to_model 인자만 지정해주면, 이를 아래와 같이 그래프 형태로 표현해준다.

 

 

 

3. add_histogram : 각 층에 있는 Weight, Bias의 분포를 볼 수 있습니다.

for name, weight in model.named_parameters():
  # 모델의 weight, bias를 로깅합니다. 'HISTOGRAMS' 탭을 클릭하여 확인할 수 있습니다.
  writer.add_histogram(tag = name, values = weight)

 

결과가 아래와 같이 나온다. 

 

처음에 layer.0과 layer.2가 무엇을 의미하는지 이해하는데 오래 걸렸는데 layer에서 trainable parameter가 있는 것을 의미한다.

 

이 가중치는 그 레이어에 있는 모든 parameter의 분포를 보여주는 거인데, 시점에 따라 다르다. 가중치 초기화를 한 이후에 바로 분포를 그리면 가중치가 어떻게 분포되어 있는지 볼 수 있고, 학습이 완료된 이후에 보면 가중치 초기화 대비 어떻게 바뀌었는지 살펴볼 수 있다.

 

 

4. add_image : 실제 이미지 데이터를 볼 수 있는 코드

 

img_grid = torchvision.utils.make_grid(images)

# Tensorboard에 이미지를 로깅합니다. 'IMAGES' 탭을 클릭하여 확인할 수 있습니다.
# matplotlib를 통해 여러 셸에서 시각화를 한다면, 한 눈에 파악하기 어려운 점이 존재합니다.
# add_image는 학습 과정에서의 이미지 변화를 트래킹하도록 설계되어 있습니다. 각 학습 global_step(epoch 혹은 step)에 대한 이미지를 저장하면, 이를 global_step에 따라 한 웹페이지에서 비교하거나 확인할 수 있습니다.
writer.add_image('MNIST Images', img_grid)

 

결과는 아래와 같이 나오는데 가로로 스와이프하면서 Step 별로 볼 수 있는 기능이 있다.

 

이를 저장하기 위해서는 아래와 같은 코드를 통해서 저장할 수 있다.

import torchvision
from torch.utils.tensorboard import SummaryWriter

# Tensorboard SummaryWriter 초기화
writer = SummaryWriter()

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 여기서 모델 학습 로직을 추가합니다.
        # 예: output = model(images), loss 계산, optimizer 업데이트 등

        if i % log_interval == 0:  # 일정 간격(log_interval)마다 로깅
            # 이미지 그리드 생성
            img_grid = torchvision.utils.make_grid(images)

            # Tensorboard에 이미지 로깅
            writer.add_image('MNIST Images', img_grid, global_step=epoch * len(train_loader) + i)

# Tensorboard 닫기
writer.close()

 

 

5. add_embedding : 고차원 데이터를 저차원 공간에 시각화하는데 사용된다. (알아서 PCA, t-SNE 같은 차원 축소 알고리즘을 사용해 뿌려준다.) 주로 단어 임베딩 단계에서 이를 비교할 때 사용한다. 이를 통해 모델이 의미상으로 어떤 것을 가깝게, 또는 멀게 해석하고 있는지 볼 수 있다.

 

# image 데이터를 embedding이라 가정하고 시각화
# 고차원 vector를 로깅합니다. 'PROJECTOR' 탭을 클릭하여 고차원 벡터를 2D, 3D 공간에 정사영한 결과물을 확인할 수 있습니다.
for images, labels in train_dataloader:
  break

embedding = images.view(images.size(0), -1)
writer.add_embedding(mat = embedding, metadata = labels, label_img = images, global_step = 1)

 

 

6. 학습 과정 logging하기

 

기존 training_loop 코드에 인자로 writer를 넣어주고 4줄만 넣어주면 자동으로 학습 과정에서 Tensorboard에 이를 저장해준다.

 

아래와 같이 저장한 지표와 정보들을 볼 수 있다.

 

 

 

 

정리해보면, 기본적으로 우리가 black-box처럼 중간중간 어떻게 모델이 구성되어 있고 어떻게 학습하고 있는지에 대해서 알 수 있다는 점이 좋은 것 같다. 다만 단순히 공부하는 입장에서는 그 효과를 파악하기가 어렵고, 실제 모델링을 하고 학습시키는 과정에서 도움이 될 것으로 보여서 실제 모델을 만들때는 이를 사용해보고 추가적으로 정리해보려고 한다.

 

 

2. WandB

- 실험 관리: WandB를 사용하면 모델의 각 학습에서 생성되는 로그를 트래킹하고 관리할 수 있다. 

- 시각화: WandB는 학습 과정에서 생성되는 metric(e.g. loss, accuracy)을 실시간으로 시각화하는 기능을 제공한다.

- 하이퍼 파라미터 최적화: WandB의 sweep 기능을 사용하면 다양한 하이퍼 파라미터 설정을 자동으로 탐색하고, 이 중에서 최고의 성능을 기록한 모델을 확인할 수 있다. 

 

추가적으로 WandB는 로컬에서 기록된 파일을 모니터링하는 TensorBoard와는 달리, 웹에서 관리되기에 여러 서버를 사용해서 학습을 진행하는 경우 하나의 웹 페이지에서 실험들을 확인할 수 있는 장점을 가진다.

 

1. 초기 세팅

#pip install wandb 실행했다는 가정.
# https://wandb.ai/authorize를 클릭하여 API key를 복사한 다음, 셸의 입력창에 붙여넣기합니다.
# `wandb.login`은 최초에 한 번만 실행하면 됩니다!

wandb.login()
run = wandb.init(project = 'test-mnist', name = 'tutorial')

#추가로 init에는 이런 인자들이 들어갈 수 있으며 이렇게 함으로써 wandb에서 다양한 모델들을 비교하고 그루핑할 수 있다.
wandb.init(project = '프로젝트명', group = '그룹명', name = '모델의 이름'
           notes = '메모(batch_size, seed, epoch 등)', 
           tags = '태그', 
           entity  = 'entity 이름')


run.finish() #WandB의 모든 남은 데이터를 서버에 업로드하고 종료한다.

 

 

 

2. wandb.log() : 로그 처리

 

위의 TensorBoard보다 훨씬 간결하게 작동하는 것을 볼 수 있다.

 

3. wandb.watch() : parameter 트래킹하기

 

run = wandb.init(project = 'test-mnist', name = 'training')

lr = 0.001
dropout_ratio = 0.2
num_epochs = 100
patience = 3
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # gpu 설정
model_name = 'exp1'
num_classes = 10

model = CNN(num_classes = 10, dropout_ratio = dropout_ratio)
model.weight_initialization()
model = model.to(device)

criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr = lr)

# wandb에 모델의 weight & bias, graident를 시각화합니다.
run.watch(model, criterion, log = 'all', log_graph = True)

model, valid_max_accuracy = training_loop(model, train_dataloader, valid_dataloader, criterion, optimizer, device, num_epochs, patience, model_name, run)
run.finish()

 

실제로 실행하면 아래와 같은 결과도 함께 나온다.

 

 

3. wandb.sweep() : 딥러닝 및 머신러닝 모델의 하이퍼 파라미터 최적화를 위해 사용되는 도구

 

코드 구성이 특이한데, 코드 동작 과정은 아래와 같이 정리해볼 수 있다.

 

1. 우선 하이퍼 파라미터 탐색에 필요한 파라미터들을 정의해줘야한다.

 

2. 하이퍼파라미터를 탐색할 모델을 정의하고 돌려주는 Run_sweep이라는 함수를 정의한다.

 

3. wandb.sweep()을 통해서 우리가 탐색하려는 Configuration, project를 정의해서 이를 sweep_id에 담는다.

• 다른 Configuration을 지정해서 다른 sweep_id에 담아서 다양하게 실험시킬 수도 있다.

 

4. 마지막으로 wandb.agent(sweep_id, run_sweep, count = 10)을 동작시킨다.

• count 인자는 10번의 다른 하이퍼파라미터 설정으로 실행하는 것.

#sweep에 필요한 configuration 설정

sweep_configuration = {
    'method': 'random',
    'metric': {'goal': 'maximize', 'name': 'valid_accuracy'},
    'parameters': {
        'lr': {'min': 0.0001, 'max': 0.01},
        'dropout_ratio': {'values': [0.1, 0.2, 0.3]},
        'weight_decay': {'min': 0.00001, 'max': 0.01}
        }
}


def run_sweep():
  num_epochs = 100
  patience = 3
  model_name = 'exp1'

  run = wandb.init(project = 'test-mnist')

  model = CNN(num_classes = 10, dropout_ratio = wandb.config.dropout_ratio)
  model.weight_initialization()
  device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  model = model.to(device)

  criterion = nn.NLLLoss()
  optimizer = optim.Adam(model.parameters(), lr = wandb.config.lr, weight_decay = wandb.config.weight_decay)

  # wandb에 모델의 weight & bias, graident를 시각화합니다.
  run.watch(model, criterion, log = 'all')
  return training_loop(model, train_dataloader, valid_dataloader, criterion, optimizer, device, num_epochs, patience, model_name, run)


sweep_id = wandb.sweep(
      sweep = sweep_configuration,
      project = 'test-mnist'
)

wandb.agent(sweep_id, function = run_sweep, count = 10)

 

 

결과로 함께 우리가 찾으려는 hyper-parameter를 함께 알려준다.

 

 

5. 여러 시각화 기능

 

글 작성 당시에 코랩이 갑자기 연결이 안되어서 wandb 사이트에서 제공하는 시각화 예시들을 가져왔다.

 

epoch에 따른 모델 예측값 추이 변화

 

 

hyper parameter 최적화에 대한 시각화

 

두 라이브러리 비교

 

1. 표면상으로는 Notebook에서 바로 조회할 수 있는 건 Tensorboard이지만, 웹에서 체계적으로 프로젝트 관리할 수 있다는 점에서는 WandB도 나름의 장점이 있어 보인다.

 

2. WandB에서는 하이퍼파라미터 조합을 찾아주도록 도움을 주는 기능을 제공하고 있으며, 시각화 관점에서 훨씬 다양하다.(아무래도 돈을 쓰고 내는 기능이다보니.. 개인 프로젝트는 무료이긴 하다.)

 

 

 

 

참고 문서

 

1. https://pytorch.org/docs/stable/tensorboard.html#torch.utils.tensorboard.writer.SummaryWriter.add_scalar  

2. https://pebpung.github.io/wandb/2021/10/06/WandB-1.html