对tensorboard进行封装
def init_tb_logger(log_dir):
from torch.utils.tensorboard import SummaryWriter
tb_logger = SummaryWriter(log_dir=log_dir)
return tb_logger
得到一个封装之后的接口 tb_logger
在训练中如何使用?
1、初始化,得到输入的地址
log_dir
2、通过获得logs文件中loss的数据,向tensorboard写入表格数据
for k, v in log_vars.items():
message += f'{k}: {v:.4e} '
# tensorboard logger
if self.use_tb_logger and 'debug' not in self.exp_name:
if k.startswith('l_'):
self.tb_logger.add_scalar(f'losses/{k}', v, current_iter)
else:
self.tb_logger.add_scalar(k, v, current_iter)
获得log_vars中的loss的名称(K)和对应的值(v)来获得写入tensorboard的数据