def val_end(self, metrics):
"""WRITE STAT FIRST"""
if self.stat_cache is None:
self.stat_cache = {'VLloss': metrics.get('VLloss'),
'VLepe': metrics.get('VLepe')}
self.savemodel({'VLloss': metrics.get('VLloss'),
'VLepe': metrics.get('VLepe')})
else:
if self.stat_cache.get('VLloss') < metrics.get('VLloss'):
self.stat_cache.update({'VLloss': metrics.get('VLloss'),
'VLepe': metrics.get('VLepe')})
self.savemodel(self.stat_cache)
else:
self.load_old_best()
self.model.eval()
with torch.no_grad():
pred_frame = metrics.get('pred_frame').cpu()
gt_frame = metrics.get('gt_frame').cpu()
pred_flow = metrics.get('pred_flow').cpu()
gt_flow = metrics.get('gt_flow').cpu()
pred_occ = replicatechannel(metrics.get('pred_occ')).cpu()
gt_occ = replicatechannel(metrics.get('gt_occ')).cpu()
data = torch.stack([pred_frame, gt_frame, pred_flow, gt_flow, pred_occ, gt_occ], 0).cpu()
data = data.reshape(-1, 3, data.size(3), data.size(4)).cpu()
grid = make_grid(data, nrow=4)
self.writer.add_images('VAL/Results', grid.unsqueeze(0), metrics.get('n_batch'))
def train_epoch_end(self, metrics):
self.model.eval()
with torch.no_grad():
pred_frame = metrics.get('pred_frame')
gt_frame = metrics.get('gt_frame')
pred_flow = metrics.get('pred_flow')
gt_flow = metrics.get('gt_flow')
pred_occ = replicatechannel(metrics.get('pred_occ'))
gt_occ = replicatechannel(metrics.get('gt_occ'))
data = torch.stack([
pred_frame.cuda(),
gt_frame.cuda(),
pred_flow.cuda(),
gt_flow.cuda(),
pred_occ.cuda(),
gt_occ.cuda()
], 0)
data = data.reshape(-1, 3, data.size(3), data.size(4)).cpu()
grid = ToTensor()((ToPILImage()(make_grid(data, nrow=2))))
self.writer.add_images('TRAIN/Results', grid.unsqueeze(0),
metrics.get('n_batch'))
self.val(metrics.get('epoch'))
def test_end(self, metrics):
self.model.eval()
with torch.no_grad():
pred_frame = metrics.get('pred_frame').cpu()
gt_frame = metrics.get('gt_frame').cpu()
pred_flow = metrics.get('pred_flow').cpu()
pred_occ = replicatechannel(metrics.get('pred_occ')).cpu()
data = torch.stack([pred_frame, gt_frame, pred_flow, pred_occ], 0)
data = data.reshape(-1, 3, data.size(3), data.size(4)).cpu()
grid = make_grid(data, nrow=4)
self.writer.add_images('Test/Results', grid.unsqueeze(0), metrics.get('n_batch'))
def train_epoch_end(self, metrics):
self.model.eval()
with torch.no_grad():
pred_frame = metrics.get('pred_frame')
gt_frame = metrics.get('gt_frame')
subf= metrics.get('subf')
pred_flow = metrics.get('pred_flow')
gt_flow = metrics.get('gt_flow')
# ff_in = metrics.get('ff_in')
# fb_in = metrics.get('fb_in')
pred_occ = replicatechannel(metrics.get('pred_occ'))
gt_occ = replicatechannel(metrics.get('gt_occ'))
data = torch.cat([pred_frame.cuda(), gt_frame.cuda(),subf.cuda(), pred_flow.cuda(), gt_flow.cuda(), pred_occ.cuda(), gt_occ.cuda()],0)
# ff_data = make_grid(torch.cat([ff_in.cuda(),fb_in.cuda()],0).cpu(), nrow=4)
data = data.cpu()
grid = make_grid(data, nrow=4)
grid = ToTensor()((ToPILImage()(grid)).resize((4106//6,2630//4)))
self.writer.add_images('TRAIN/Results', grid.unsqueeze(0), metrics.get('n_batch'))
# self.writer.add_images('INPUT/Flows_prior', ff_data.unsqueeze(0), metrics.get('n_batch'))
self.val(metrics.get('epoch'))