TensorBoard in PyTorch

TensorBoard in PyTorch

python -m pip install torch tensorboard torch-tb-profiler

import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter

# log_dir : Log File을 저장할 Directory, Default : ./runs/CURRENT_DATETIME_HOSTNAME
# comment : Log Directory 이름에 추가되는 주석, Default : '' (Empty String)
writer = SummaryWriter(log_dir="./logs", comment="First Exp")

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

# Define Model, Criterion, Optimizer
model = torch.nn.Linear(1, 1)
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# Example Image and Label
image_batch = torch.randn(4, 1, 28, 28)  # Image [Batch, Channel, Height, Width]
label_batch = torch.randint(0, 2, (4, 28, 28))  # Label (Binary Classification)

def train_model(iter):
    for epoch in range(iter):
        # Forward Pass
        y1 = model(x)
        loss = criterion(y1, y)
        
        # add_scalar([Graph Title], [y_axis_value], [x_axis_value])
        writer.add_scalar("Train/loss", loss, global_step=epoch)  # Write Log Data

        # Backward Pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # Visualize sample images, predictions and labels every 5 epochs
        if epoch % 5 == 0:
            # Normalize the image to [0, 1] range for visualization
            image = image_batch[1, 0:1, :, :]
            image = (image - image.min()) / (image.max() - image.min())
            # Images
            writer.add_image('train/Image', image, epoch)
            
            # Predictions
            output_labels = torch.argmax(F.softmax(y1, dim=1), dim=1, keepdim=True)
            writer.add_image('train/Prediction', output_labels[0].unsqueeze(0) * 50, epoch)
            
            # Labels
            labels = label_batch[1, ...].unsqueeze(0) * 50
            writer.add_image('train/GroundTruth', labels, epoch)

train_model(10)

writer.flush()  # Write to Disk
writer.close()  # Close Writer

Remote Server의 TensorBoard를 Local에서 확인

1. 명령어 입력 방법

• 1-1. Remote Server에서 아래 명령어 실행

# tensorboard --logdir [Log File이 저장된 Directory]

tensorboard --logdir ./logs

• 1-2. [1-1]의 명령어를 입력하였을 때 나오는 주소에 접속
  

2. Visual Studio Code의 내장 기능 활용 방법

• 2-1. TensorBoard Library를 불러오는 곳에 위치한 [Launch TensorBoard Session] 실행