torch中使用tensorboard记录loss

首先建立一个.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后记录