pytorch tensorboardX获取writer实例

Great_lid

已于 2022-12-14 20:07:06 修改

阅读量166

点赞数

文章标签： pytorch 人工智能 python

于 2022-12-14 19:59:16 首次发布

本文链接：https://blog.csdn.net/pp2373886592/article/details/128321191

版权

自定义函数，获取tensorboardX里的writer实例并使用其记录数据

import os
from datetime import datetime
import socket
from tensorboardX import SummaryWriter

def get_summary_writer(log_dir, prefix=None):
    if not os.path.exists(log_dir):
        os.mkdir(log_dir)
    if prefix is None:
        log_dir = os.path.join(log_dir, datetime.now().strftime('%b%d_%H-%M-%S')+'_'+socket.gethostname())
    else:
        log_dir = os.path.join(log_dir, prefix+'_'+datetime.now().strftime('%b%d_%H-%M-%S')+'_'+socket.gethostname())
    if not os.path.exists(log_dir):
        os.mkdir(log_dir)
    writer = SummaryWriter(log_dir)
    return writer

writer = SummaryWriter(log_dir)
writer.add_scalar(os.path.join(prefix, 'train_loss'), loss_data, iteration)