[PyTorch] Tensorboard 사용하기 (colab)

텐서보드는 딥러닝 학습시 시각화를 돕는 도구이다. colab 환경에서Tensorboard를 구현 하는 방법을 살펴보겠다.

공식 문서: https://tutorials.pytorch.kr/recipes/recipes/tensorboard_with_pytorch.html

import torch
from tensorboardX import SummaryWriter
writer = SummaryWriter()

먼저 tensorboardX의 SummaryWriter()를 불러온다.

---

x = torch.arange(-5, 5, 0.1).view(-1, 1)
y = -5 * x + 0.1 * torch.randn(x.size())

model = torch.nn.Linear(1, 1)
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr = 0.1)

def train_model(iter):
    for epoch in range(iter):
        y1 = model(x)
        loss = criterion(y1, y)
        writer.add_scalar("Loss/train", loss, epoch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

train_model(10)
writer.flush()

간단한 linear모델을 구현해서 loss를 시각화 해보도록 하겠다.

 

loss와 epoch같은 스칼라 값을 기록할 때 add_scalar를 사용한다.

• add_scalar(tag, scalar_value, global_step=None, walltime=None)

---

writer.close()

%load_ext tensorboard
%tensorboard --logdir=runs

colab 환경에서 구현하기 위해 magic command(%) 이용한다. 아래와 같이 Tensorboard가 로드된다.