Pytorch+Tensorboard使用

Pytorch+Tensorboard使用

Tensorboard在神经网络训练时能够实时跟踪训练结果，查看loss，accuracy等，以下介绍使用方法

示例程序

from torch.utils.tensorboard import SummaryWriter
from Test3_vggnet.model import vgg #随便一个pytorch写的网络都可以
import torch
import time
import os


cur_dir = os.path.dirname(__file__)
log_dir = os.path.join(cur_dir,"logs/tensorboard")
writer = SummaryWriter(log_dir=log_dir)

# show the net structure
net = vgg()
dummy_input = torch.rand(13, 3, 224, 224) # 模拟输出
writer.add_graph(net, dummy_input)

# 模拟训练过程
loss = 5005*100
accuracy = 0.
for epoch in range(100):
    for step, value in enumerate(range(5005)):
        time.sleep(0.001)
        loss -= 1
        accuracy += 1/(5005*100)
        
        if step % 1000 == 999:
            
            writer.add_scalars("loss", {"train" : loss,
                                        "val" : 5005*100-loss}, epoch*5005+step)
            
            writer.add_scalars("accuracy", {"train" : accuracy,
                                            "val" : 1-accuracy}, epoch*5005+step)

print("finished !!")

该程序首先在tensorboard上展示了网络结构，然后显示实时的loss以及accuracy，在一个scalars里面的内容会在一张图上显示

网络结构，可以双击进入查看

loss结果