#两种import方式. 
# 1 :
# from torch.utils.tensorboard import SummaryWriter
# 2 :
from tensorboardX import SummaryWriter

import torch.optim
from torch.optim.lr_scheduler import StepLR

now = time.strftime('%Y%m%d-%H_%M_%S', time.localtime(time.time()))
writer = SummaryWriter(log_dir=f'{LOG_DIR}/{now}')

# set optimizer 设置优化器
optimizer = torch.optim.SGD(...)
# set scheduler 设置学习率调整策略
scheduler = StepLR(optimizer, ...)

...
def train(..., epoch):
    ...
    for i, (input, target) in enumerate(train_loader):
      ...
	
  
    writer.add_scalars(main_tag='Loss',tag_scalar_dict={
        'train loss' : loss,
    }, global_step=epoch)

    writer.add_scalars(main_tag='Prec',tag_scalar_dict={
        'train top1' : top1,
        'trian top5' : top5
    }, global_step=epoch)

    writer.add_scalars(main_tag='time/train', tag_scalar_dict={
        'batch_time': batch_time,
        'data_time': data_time,
        },
        global_step=epoch)

def validate(..., epoch):
    ...
    with torch.no_grad():
        ...

      writer.add_scalars(main_tag='Loss', tag_scalar_dict={
          'val loss' : losses.avg,
      }, global_step=epoch)

      writer.add_scalars(main_tag='Prec', tag_scalar_dict={
          'validate top1' : top1.avg,
          'validate top5' : top5.avg
      }, global_step=epoch)

      writer.add_scalar('time/val/batch_time', batch_time.avg, epoch)

    return ...

def main():
    ...
    for epoch in range(args.start_epoch, args.epochs):
        ...
        train(..., epoch)
        validate(..., epoch)

        writer.add_scalar('learning rate', optimizer.state_dict()['param_groups'][0]['lr'], epoch)
        scheduler.step()

        ...
    
    writer.close()
    # only close once after all add_ function is done
    # the end of main()