[PyTorch] 학습 Loss 그래프 그리기 (Tensorboard 이용)

 

모델 학습 과정에서의 Loss를 그래프로 시각화하기 위해, Tensorboard를 이용하는 방법을 알아본다.

 

참조1 : seongkyun.github.io/others/2019/05/11/pytorch_tensorboard/

참조2 : sensibilityit.tistory.com/512

참조한 블로그들이다. 이 포스팅에서는 핵심만 간략하게 설명할 것이니 자세한 설명이 필요하다면 위 블로그에서 확인하면 된다 :)

 

 

1. Tensorboard 설치하기

!pip install tensorboardX

tensorflow도 설치되어 있어야한다고 한다. 

 

 

2. 모듈 임포트 및 객체 생성

from tensorboardX import SummaryWriter
summary = SummaryWriter()

 

 

3. loss 추가하기

이제 모델 학습 코드 부분을 보면 학습을 여러 번 진행되도록 하는 반복문이 있을 것이다. 여기서 반복 횟수를 의미하는 변수명을 체크한다.

나의 경우 변수명은 iteration

 

반복문 내 최하단에 다음과 같은 코드를 추가한다.

if iteration % 10 == 0:    # 매 10 iteration마다 업데이트
	writer.add_scalar('D_loss_adv', d_loss_adv.item(), iteration)
	writer.add_scalar('D_loss_cls', d_loss_cls.item(), iteration)

보통 10~20 iteration마다 업데이트를 한다고 한다.

add_scalar()의 인자에는 loss명을 설정해 넘겨주고, 그래프에 출력할 loss를 .item()으로 가져와 넘겨주고, 마지막 인자에는 반복횟수를 의미하는 변수명을 넘겨준다.

 

 

4. 모델 학습 진행

 

 

5. Loss 그래프 출력하기

모델 학습이 완료되었으면, 다음 코드를 터미널에서 실행한다.

그러면 TensorBoard 창에서 그래프가 출력된다.

 

Google Colab 환경인 경우 아래 코드를 입력한다.

# Colab
%load_ext tensorboard
%tensorboard --logdir runs --port=6006

PyCharm과 같은 에디터 환경인 경우 아래 코드를 입력한다.

# 에디터
tensorboard --logdir runs