| | import torch
|
| |
|
| | from utils.flow_viz import flow_tensor_to_image
|
| |
|
| |
|
| | class Logger:
|
| | def __init__(self, lr_scheduler,
|
| | summary_writer,
|
| | summary_freq=100,
|
| | start_step=0,
|
| | ):
|
| | self.lr_scheduler = lr_scheduler
|
| | self.total_steps = start_step
|
| | self.running_loss = {}
|
| | self.summary_writer = summary_writer
|
| | self.summary_freq = summary_freq
|
| |
|
| | def print_training_status(self, mode='train'):
|
| |
|
| | print('step: %06d \t epe: %.3f' % (self.total_steps, self.running_loss['epe'] / self.summary_freq))
|
| |
|
| | for k in self.running_loss:
|
| | self.summary_writer.add_scalar(mode + '/' + k,
|
| | self.running_loss[k] / self.summary_freq, self.total_steps)
|
| | self.running_loss[k] = 0.0
|
| |
|
| | def lr_summary(self):
|
| | lr = self.lr_scheduler.get_last_lr()[0]
|
| | self.summary_writer.add_scalar('lr', lr, self.total_steps)
|
| |
|
| | def add_image_summary(self, img1, img2, flow_preds, flow_gt, mode='train',
|
| | ):
|
| | if self.total_steps % self.summary_freq == 0:
|
| | img_concat = torch.cat((img1[0].detach().cpu(), img2[0].detach().cpu()), dim=-1)
|
| | img_concat = img_concat.type(torch.uint8)
|
| |
|
| | flow_pred = flow_tensor_to_image(flow_preds[-1][0])
|
| | forward_flow_gt = flow_tensor_to_image(flow_gt[0])
|
| | flow_concat = torch.cat((torch.from_numpy(flow_pred),
|
| | torch.from_numpy(forward_flow_gt)), dim=-1)
|
| |
|
| | concat = torch.cat((img_concat, flow_concat), dim=-2)
|
| |
|
| | self.summary_writer.add_image(mode + '/img_pred_gt', concat, self.total_steps)
|
| |
|
| | def push(self, metrics, mode='train'):
|
| | self.total_steps += 1
|
| |
|
| | self.lr_summary()
|
| |
|
| | for key in metrics:
|
| | if key not in self.running_loss:
|
| | self.running_loss[key] = 0.0
|
| |
|
| | self.running_loss[key] += metrics[key]
|
| |
|
| | if self.total_steps % self.summary_freq == 0:
|
| | self.print_training_status(mode)
|
| | self.running_loss = {}
|
| |
|
| | def write_dict(self, results):
|
| | for key in results:
|
| | tag = key.split('_')[0]
|
| | tag = tag + '/' + key
|
| | self.summary_writer.add_scalar(tag, results[key], self.total_steps)
|
| |
|
| | def close(self):
|
| | self.summary_writer.close()
|
| |
|
