pytorch学习系列(3)：两种tensorboard可视化

1.tensorboardX可视化

这种方式直接使用tensorboardX中的SummaryWriter来保存变量。

from tensorboardX import SummaryWriter
writer = SummaryWriter(log_dir=log_dir)#log_dir为日志文件的保存目录
train()
dummy_input = torch.rand(opt.batch_size,channels,opt.image_size,opt.image_size)#设定一个假的输入量
model = net.NET()
writer.add_graph(net,(dummy_input,))#保存计算图，可视化模型结构
writer.add_scalar('train/loss', your_loss, epoch)#保存标量，一般为损失或者精度
x = vutils.make_grid(image,normalize=True,scale_each=True)
writer.add_image('image',x,epoch)#可视化中间过程的图像
for name,param in model.named_parameters():
            writer.add_histogram(name,param.clone().cpu().data.numpy(),epoch)#保存参数的直方图

2.利用脚本实现tensorboard

复制如下代码到新的文件logger.py中，将其放入项目目录中。

import tensorflow as tf
import numpy as np
import scipy.misc
try:
    from StringIO import StringIO  # Python 2.7
except ImportError:
    from io import BytesIO         # Python 3.x


class Logger(object):
    
    def __init__(self, log_dir):
        """Create a summary writer logging to log_dir."""
        self.writer = tf.summary.FileWriter(log_dir)

    def scalar_summary(self, tag, value, step):
        """Log a scalar variable."""
        summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
        self.writer.add_summary(summary, step)

    def image_summary(self, tag, images, step):
        """Log a list of images."""

        img_summaries = []
        for i, img in enumerate(images):
            # Write the image to a string
            try:
                s = StringIO()
            except:
                s = BytesIO()
            scipy.misc.toimage(img).save(s, format="png")

            # Create an Image object
            img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(),
                                       height=img.shape[0],
                                       width=img.shape[1])
            # Create a Summary value
            img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum))

        # Create and write Summary
        summary = tf.Summary(value=img_summaries)
        self.writer.add_summary(summary, step)
        
    def histo_summary(self, tag, values, step, bins=1000):
        """Log a histogram of the tensor of values."""

        # Create a histogram using numpy
        counts, bin_edges = np.histogram(values, bins=bins)

        # Fill the fields of the histogram proto
        hist = tf.HistogramProto()
        hist.min = float(np.min(values))
        hist.max = float(np.max(values))
        hist.num = int(np.prod(values.shape))
        hist.sum = float(np.sum(values))
        hist.sum_squares = float(np.sum(values**2))

        # Drop the start of the first bin
        bin_edges = bin_edges[1:]

        # Add bin edges and counts
        for edge in bin_edges:
            hist.bucket_limit.append(edge)
        for c in counts:
            hist.bucket.append(c)

        # Create and write Summary
        summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)])
        self.writer.add_summary(summary, step)
        self.writer.flush()

在自己的train程序中导入logger.py

from logger import Logger
logger = Logger(log_dir)#log_dir为日志文件的保存目录

在训练过程中以如下方式记录

if idx % 20== 0:
    # 1. Log scalar values (scalar summary)
    # 日志输出标量信息（scalar summary）
    info = { 'loss': mse.item()}#损失函数

    for tag, value in info.items():
        logger.scalar_summary(tag, value, idx)

    # 2. Log values and gradients of the parameters (histogram summary)
    # 日志输出参数值和梯度（histogram summary)
    for tag, value in net.named_parameters():
        tag = tag.replace('.', '/')
        logger.histo_summary(tag, value.data.cpu().numpy(), idx)
        logger.histo_summary(tag+'/grad', value.grad.data.cpu().numpy(),idx)

    # 3. Log training images (image summary)
    # 日志输出图像(image summary)
    info = { 'recon': output.view(-1, opt.image_size, opt.image_size)[:5].cpu().detach().numpy(),
            'orign': input.view(-1, opt.image_size, opt.image_size)[:5].cpu().detach().numpy() }
            #由于我做的是图像压缩，所以这里有recon和orign两个，[:5]表示显示五个图像

    for tag, images in info.items():
        logger.image_summary(tag, images, idx)

我认为还是第一种方式更简单明了一些。