首先建立一个.py文件叫做utils.py 文件内容如下:
# -*- coding: utf-8 -*-
from tensorboardX import SummaryWriter
class TensorboardWriter(SummaryWriter):
def __init__(self, logdir):
super(TensorboardWriter, self).__init__(logdir)
self.logdir = self.file_writer.get_logdir()
def update_parameters(self, module, step_i):
"""
module: nn.Module
"""
for name, param in module.named_parameters():
self.add_histogram(name, param.clone().cpu().data.numpy(), step_i)
def update_loss(self, loss, step_i, name='loss'):
self.add_scalar(name, loss, step_i)
def update_histogram(self, values, step_i, name='hist'):
self.add_histogram(name, values, step_i)
然后在train.py文件中导入
from utils import TensorboardWriter
在文中进行初始化
writer = TensorboardWriter(args.log_dir)
产生loss后记录