Tensorboard를 활용한 Loss 시각화 (Pytorch)

텐서보드는 머신러닝 실험에 필요한 다양한 시각화 및 도구를 제공하는 구글에서 개발한 오픈소스입니다. 해당 문서에서는 pytorch에서 tensorboard를 활용하여 학습 진행 과정에 따른 loss 변화를 시각화하는 방법을 기술합니다.

1. tensorboard 모듈 설치

$ pip install tensorboard

2. 필수 모듈 import하기

from torch.utils.tensorboard import SummaryWriter

3. tensorboard에 값을 기록하기 위한 SummaryWriter 객체 생성

writer = SummaryWriter('scalar/')

현재 디렉토리에 scalar 폴더가 생성되며 해당 폴더 내부에 loss값을 기록한 로그 파일이 저장됩니다.

그 다음, train 함수의 forward 연산 중 loss를 업데이트하는 부분에서 add_scalar 함수를 통해 기록하고자 하는 변수를 parameter로 넘겨줍니다.

for it, (x, y) in pbar:
# forward the model
	with torch.set_grad_enabled(is_train):
 		logits, loss = model(x, y)
		loss = loss.mean() 
		writer.add_scalar("Loss/train", loss, epoch)
		losses.append(loss.item())

학습이 종료된 후에는 writer 객체를 close하여 값을 저장합니다. close하지 않으면 loss 값이 저장되지 않으니 주의합니다.

writer.close() ## 학습이 종료된 후에 반드시 선언할 것!

scalar 폴더 내부에 다음과 같이 event log 파일이 생성되었으면 성공적으로 loss가 기록된 것입니다.

프로그램이 종료되었으면, 다음 명령어를 실행하며 브라우저를 연 뒤 tensorboard에 기록된 값을 확인합니다.

$ tensorboard --logdir ./scalar

4. 실행 결과