tensorboard 记录 Loss 和评估指标曲线
- import torch
- from torch import nn
- torch.backends.cudnn.benchmark = True
- from torch.utils.tensorboard import SummaryWriter
- ######### Set Seeds ###########
- random.seed(1234)
- np.random.seed(1234)
- torch.manual_seed(1234)
- torch.cuda.manual_seed_all(1234)
- def train():
- # --------------------------------------------------------------
- time_stamp = "{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.datetime.now())
- comment = f'bath_size={config.BATCH_SIZE}_lr={config.LR}'
- writer = SummaryWriter(log_dir="workdir/" + time_stamp, comment=comment)
- # --------------------------------------------------------------
- for epoch in range(start_epoch, EPOCH):
- net.train()
- for step, (x, y, path_code) in enumerate((train_loader), 0):
- for index in range(1, 2):
- tmp = torch.zeros(size=(x.size(0), 3, 3, x.size(3), x.size(4)))
- tmp[:, :, :, :, :] = x[:, index - 1:index + 2, :, :, :]
- tmp = tmp.cuda()
- tmp_y = torch.zeros(size=(y.size(0), 3, y.size(3), y.size(4)))
- tmp_y[:, :, :, :, ] = y[:, index, :, :, :]
- tmp_y = tmp_y.cuda()
- y_hat, output, target = net(tmp)
- losses_moco += constrastlossprint.item()
- losses += ssimlossprint.item() # + vgglossprint.item()
- ### overall loss
- overall_loss = SSIM_LOSS + MOCO_LOSS
- optimizer.zero_grad()
- ### backward loss
- overall_loss.backward()
- ### update parameters
- optimizer.step()
- # profiler.step()
- losses /= (index) # 一个视频中的一帧的平均 loss
- losses_moco /= index
- # print('losses_pre_%d_video:' % step, losses)
- losses_epoch += losses # 所有视频的平均 loss 的和
- losses_moco_epoch += losses_moco
- writer.add_scalar('SSIM Loss', losses_epoch, global_step=epoch)
- writer.add_scalar('Moco Loss', losses_moco_epoch, global_step=epoch)
- scheduler.step()
- if __name__ == '__main__':
- train()
torch.profiler.profile 的使用 (pytorch1.8 以上)
- with torch.profiler.profile(
- activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
- schedule=torch.profiler.schedule(
- wait=2,
- warmup=2,
- active=6,
- repeat=1),
- on_trace_ready=torch.profiler.tensorboard_trace_handler('./workdir'),
- # with_trace=True
- ) as profiler:
- for step, (x, y, path_code) in enumerate((train_loader), 0):
- xxxxxxx
- ### overall loss
- overall_loss = SSIM_LOSS + MOCO_LOSS
- optimizer.zero_grad()
- ### backward loss
- overall_loss.backward()
- ### update parameters
- optimizer.step()
- profiler.step() # 注意 *******
来源: http://www.bubuko.com/infodetail-3815029.html