def mini_eval(self, current_step):
if current_step % self.config.TRAIN.mini_eval_step != 0:
return
temp_meter_frame = AverageMeter()
temp_meter_flow = AverageMeter()
self.G.eval()
self.D.eval()
for data in self.val_dataloader:
# get the data
target_mini = data[:, :, 1, :, :]
input_data_mini = data[:, :, 0, :, :]
# squeeze the dimension
target_mini = target_mini.view(target_mini.shape[0], -1,
target_mini.shape[-2],
target_mini.shape[-1]).cuda()
input_data_mini = input_data_mini.view(
input_data_mini.shape[0], -1, input_data_mini.shape[-2],
input_data_mini.shape[-1]).cuda()
# Use the model, get the output
output_flow_G_mini, output_frame_G_mini = self.G(input_data_mini)
input_gtFlowEstimTensor = torch.cat([input_data_mini, target_mini],
1)
gtFlow, _ = flow_batch_estimate(self.F, input_gtFlowEstimTensor)
frame_psnr_mini = psnr_error(output_frame_G_mini.detach(),
target_mini)
flow_psnr_mini = psnr_error(output_flow_G_mini.detach(), gtFlow)
temp_meter_frame.update(frame_psnr_mini.detach())
temp_meter_flow.update(flow_psnr_mini.detach())
self.logger.info(
f'&^*_*^& ==> Step:{current_step}/{self.max_steps} the frame PSNR is {temp_meter_frame.avg:.3f}, the flow PSNR is {temp_meter_flow.avg:.3f}'
)
def mini_eval(self, current_step):
if current_step % 10 != 0 or current_step == 0:
return
temp_meter_frame = AverageMeter()
temp_meter_flow = AverageMeter()
self.G.eval()
self.D.eval()
for data in self.val_dataloader:
# base on the D to get each frame
target_mini = data[:, :, -1, :, :].cuda() # t+1 frame
input_data = data[:, :, :-1, :, :] # 0 ~ t frame
input_last_mini = input_data[:, :, -1, :, :].cuda() # t frame
# squeeze the D dimension to C dimension, shape comes to [N, C, H, W]
input_data_mini = input_data.view(input_data.shape[0], -1, input_data.shape[-2], input_data.shape[-1]).cuda()
output_pred_G = self.G(input_data_mini)
gtFlow, _ = flow_batch_estimate(self.F, torch.cat([input_last_mini, target_mini], 1))
predFlow, _ = flow_batch_estimate(self.F, torch.cat([input_last_mini, output_pred_G], 1))
frame_psnr_mini = psnr_error(output_pred_G.detach(), target_mini)
flow_psnr_mini = psnr_error(predFlow, gtFlow)
temp_meter_frame.update(frame_psnr_mini.detach())
temp_meter_flow.update(flow_psnr_mini.detach())
self.logger.info(f'&^*_*^& ==> Step:{current_step}/{self.max_steps} the PSNR is {temp_meter_frame.avg:.2f}, the flow PSNR is {temp_meter_flow.avg:.2f}')
def train(self,current_step):
# Pytorch [N, C, D, H, W]
# initialize
start = time.time()
self.G.train()
self.D.train()
writer = self.kwargs['writer_dict']['writer']
global_steps = self.kwargs['writer_dict']['global_steps_{}'.format(self.kwargs['model_type'])]
# get the data
data = next(self._train_loader_iter)
self.data_time.update(time.time() - start)
# base on the D to get each frame
target = data[:, :, -1, :, :].cuda() # t+1 frame
input_data = data[:, :, :-1, :, :] # 0 ~ t frame
input_last = input_data[:, :, -1, :, :].cuda() # t frame
# squeeze the D dimension to C dimension, shape comes to [N, C, H, W]
input_data = input_data.view(input_data.shape[0], -1, input_data.shape[-2], input_data.shape[-1]).cuda()
# True Process =================Start===================
#---------update optim_G ---------
self.set_requires_grad(self.D, False)
output_pred_G = self.G(input_data)
predFlowEstim = torch.cat([input_last, output_pred_G],1)
gtFlowEstim = torch.cat([input_last, target], 1)
_, gtFlow = flow_batch_estimate(self.F, gtFlowEstim)
_, predFlow = flow_batch_estimate(self.F, predFlowEstim)
# loss_g_adv = self.g_adv_loss(self.D(G_output))
loss_g_adv = self.gan_loss(self.D(output_pred_G), True)
loss_op = self.op_loss(predFlow, gtFlow)
loss_int = self.int_loss(output_pred_G, target)
loss_gd = self.gd_loss(output_pred_G, target)
loss_g_all = self.loss_lamada['intentsity_loss'] * loss_int + self.loss_lamada['gradient_loss'] * loss_gd + self.loss_lamada['opticalflow_loss'] * loss_op + self.loss_lamada['gan_loss'] * loss_g_adv
self.optim_G.zero_grad()
loss_g_all.backward()
self.optim_G.step()
# record
self.loss_meter_G.update(loss_g_all.detach())
if self.config.TRAIN.adversarial.scheduler.use:
self.lr_g.step()
#---------update optim_D ---------------
self.set_requires_grad(self.D, True)
self.optim_D.zero_grad()
# G_output = self.G(input)
temp_t = self.D(target)
temp_g = self.D(output_pred_G.detach())
loss_d_1 = self.gan_loss(temp_t, True)
loss_d_2 = self.gan_loss(temp_g, False)
loss_d = (loss_d_1 + loss_d_2) * 0.5
# loss_d.sum().backward()
loss_d.backward()
self.optim_D.step()
if self.config.TRAIN.adversarial.scheduler.use:
self.lr_d.step()
self.loss_meter_D.update(loss_d.detach())
# ======================End==================
self.batch_time.update(time.time() - start)
if (current_step % self.log_step == 0):
msg = 'Step: [{0}/{1}]\t' \
'Type: {cae_type}\t' \
'Time: {batch_time.val:.2f}s ({batch_time.avg:.2f}s)\t' \
'Speed: {speed:.1f} samples/s\t' \
'Data: {data_time.val:.2f}s ({data_time.avg:.2f}s)\t' \
'Loss_G: {losses_G.val:.5f} ({losses_G.avg:.5f})\t' \
'Loss_D:{losses_D.val:.5f}({losses_D.avg:.5f})'.format(current_step, self.max_steps, cae_type=self.kwargs['model_type'], batch_time=self.batch_time, speed=self.config.TRAIN.batch_size/self.batch_time.val, data_time=self.data_time,losses_G=self.loss_meter_G, losses_D=self.loss_meter_D)
self.logger.info(msg)
writer.add_scalar('Train_loss_G', self.loss_meter_G.val, global_steps)
writer.add_scalar('Train_loss_D', self.loss_meter_D.val, global_steps)
if (current_step % self.vis_step == 0):
vis_objects = OrderedDict()
vis_objects['train_target'] = target.detach()
vis_objects['train_output_pred_G'] = output_pred_G.detach()
vis_objects['train_gtFlow'] = gtFlow.detach()
vis_objects['train_predFlow'] = predFlow.detach()
training_vis_images(vis_objects, writer, global_steps)
global_steps += 1
# reset start
start = time.time()
self.saved_model = {'G':self.G, 'D':self.D}
self.saved_optimizer = {'optim_G': self.optim_G, 'optim_D': self.optim_D}
self.saved_loss = {'loss_G':self.loss_meter_G.val, 'loss_D':self.loss_meter_D.val}
self.kwargs['writer_dict']['global_steps_{}'.format(self.kwargs['model_type'])] = global_steps
def evaluate(self, current_step):
'''
Evaluate the results of the model
!!! Will change, e.g. accuracy, mAP.....
!!! Or can call other methods written by the official
'''
tb_writer = self.trainer.kwargs['writer_dict']['writer']
global_steps = self.trainer.kwargs['writer_dict'][
'global_steps_{}'.format(self.trainer.kwargs['model_type'])]
frame_num = self.trainer.config.DATASET.test_clip_length
psnr_records = []
score_records = []
# score_records_w=[]
w_dict = OrderedDict()
# total = 0
# calc the scores for the training set
for video_name in self.trainer.test_dataset_keys_w:
dataset = self.trainer.test_dataset_dict_w[video_name]
len_dataset = dataset.pics_len
test_iters = len_dataset - frame_num + 1
test_counter = 0
data_loader = DataLoader(dataset=dataset,
batch_size=1,
shuffle=False,
num_workers=1)
scores = [0.0 for i in range(len_dataset)]
for test_input in data_loader:
# import ipdb; ipdb.set_trace()
test_target = test_input[:, -1].cuda()
test_input = test_input[:, :-1].view(
test_input.shape[0], -1, test_input.shape[-2],
test_input.shape[-1]).cuda()
# import ipdb; ipdb.set_trace()
g_output_flow, g_output_frame = self.trainer.G(test_input)
gt_flow_esti_tensor = torch.cat([test_input, test_target], 1)
# flow_gt, _ = self.trainer.batch_estimate(gt_flow_esti_tensor)
flow_gt, _ = flow_batch_estimate(self.trainer.F,
gt_flow_esti_tensor)
# score = amc_score(test_input, g_output_frame, flow_gt, g_output_flow)
diff_appe, diff_flow = simple_diff(test_input, g_output_frame,
flow_gt, g_output_flow)
max_patch_appe, max_patch_flow = find_max_patch(
diff_appe, diff_flow)
score_appe = torch.mean(max_patch_appe)
score_flow = torch.mean(max_patch_flow)
scores[test_counter + frame_num - 1] = [score_appe, score_flow]
test_counter += 1
if test_counter >= test_iters:
scores[:frame_num - 1] = [scores[frame_num - 1]]
scores = torch.tensor(scores)
frame_w = torch.mean(scores[:, 0])
flow_w = torch.mean(scores[:, 1])
w_dict[video_name] = [len_dataset, frame_w, flow_w]
print(
f'finish calc the scores of training set {video_name} in step:{current_step}'
)
break
wf, wi = calc_w(w_dict)
# wf , wi = 1.0, 1.0
tb_writer.add_text('weight of train set',
f'w_f:{wf:.3f}, w_i:{wi:.3f}', global_steps)
num_videos = 0
random_video_sn = torch.randint(0, len(self.trainer.test_dataset_keys),
(1, ))
# calc the score for the test dataset
for sn, video_name in enumerate(self.trainer.test_dataset_keys):
num_videos += 1
# need to improve
dataset = self.trainer.test_dataset_dict[video_name]
len_dataset = dataset.pics_len
test_iters = len_dataset - frame_num + 1
test_counter = 0
data_loader = DataLoader(dataset=dataset,
batch_size=1,
shuffle=False,
num_workers=1)
vis_range = range(int(len_dataset * 0.5),
int(len_dataset * 0.5 + 5))
psnrs = np.empty(shape=(len_dataset, ), dtype=np.float32)
scores = np.empty(shape=(len_dataset, ), dtype=np.float32)
# scores = [0.0 for i in range(len_dataset)]
for frame_sn, test_input in enumerate(data_loader):
test_target = test_input[:, -1].cuda()
test_input = test_input[:, :-1].view(
test_input.shape[0], -1, test_input.shape[-2],
test_input.shape[-1]).cuda()
g_output_flow, g_output_frame = self.trainer.G(test_input)
gt_flow_esti_tensor = torch.cat([test_input, test_target], 1)
# flow_gt,_ = self.trainer.batch_estimate(gt_flow_esti_tensor)
flow_gt, _ = flow_batch_estimate(self.trainer.F,
gt_flow_esti_tensor)
# score = amc_score(test_input, g_output_frame, flow_gt, g_output_flow)
test_psnr = psnr_error(g_output_frame, test_target)
score = amc_score(test_input, g_output_frame, flow_gt,
g_output_flow, wf, wi)
# import ipdb; ipdb.set_trace()
test_psnr = test_psnr.tolist()
score = score.tolist()
psnrs[test_counter + frame_num - 1] = test_psnr
scores[test_counter + frame_num - 1] = score
test_counter += 1
if sn == random_video_sn and (frame_sn in vis_range):
self.add_images(test_target, flow_gt, g_output_frame,
g_output_flow, tb_writer, global_steps)
if test_counter >= test_iters:
psnrs[:frame_num - 1] = psnrs[frame_num - 1]
# import ipdb; ipdb.set_trace()
scores[:frame_num - 1] = (
scores[frame_num - 1],
) # fix the bug: TypeError: can only assign an iterable
smax = max(scores)
normal_scores = np.array(
[np.divide(s, smax) for s in scores])
psnr_records.append(psnrs)
score_records.append(normal_scores)
print(f'finish test video set {video_name}')
break
result_dict = {
'dataset': self.trainer.config.DATASET.name,
'psnr': psnr_records,
'flow': [],
'names': [],
'diff_mask': [],
'score': score_records,
'num_videos': num_videos
}
result_path = os.path.join(
self.trainer.config.TEST.result_output,
f'{self.trainer.verbose}_cfg#{self.trainer.config_name}#step{current_step}@{self.trainer.kwargs["time_stamp"]}_results.pkl'
)
with open(result_path, 'wb') as writer:
pickle.dump(result_dict, writer, pickle.HIGHEST_PROTOCOL)
# results = eval_api.evaluate('compute_auc_score', result_path, self.trainer.logger, self.trainer.config)
results = self.trainer.evaluate_function(
result_path, self.trainer.logger, self.trainer.config,
self.trainer.config.DATASET.score_type)
self.trainer.logger.info(results)
tb_writer.add_text('amc: AUC of ROC curve', f'auc is {results.auc}',
global_steps)
return results.auc
def train(self, current_step):
# Pytorch [N, C, D, H, W]
# initialize
start = time.time()
self.G.train()
self.D.train()
writer = self.kwargs['writer_dict']['writer']
global_steps = self.kwargs['writer_dict']['global_steps_{}'.format(
self.kwargs['model_type'])]
# get the data
data = next(self._train_loader_iter)
self.data_time.update(time.time() - start)
# base on the D to get each frame
# in this method, D = 2 and not change
target = data[:, :, 1, :, :] # target(2-nd) frame
input_data = data[:, :, 0, :, :] # input(1-st) frame
# squeeze the D dimension to C dimension, shape comes to [N, C, H, W]
target = target.reshape(target.shape[0], -1, target.shape[-2],
target.shape[-1]).cuda()
input_data = input_data.reshape(input_data.shape[0], -1,
input_data.shape[-2],
input_data.shape[-1]).cuda()
# True Process =================Start===================
#---------update optim_G ---------
self.set_requires_grad(self.D, False)
output_flow_G, output_frame_G = self.G(input_data)
gt_flow_esti_tensor = torch.cat([input_data, target], 1)
flow_gt, _ = flow_batch_estimate(
self.F,
gt_flow_esti_tensor,
normalize=self.config.ARGUMENT.train.normal.use,
mean=self.config.ARGUMENT.train.normal.mean,
std=self.config.ARGUMENT.train.normal.mean)
fake_g = self.D(torch.cat([target, output_flow_G], dim=1))
loss_g_adv = self.gan_loss(fake_g, True)
loss_op = self.op_loss(output_flow_G, flow_gt)
loss_int = self.int_loss(output_frame_G, target)
loss_gd = self.gd_loss(output_frame_G, target)
loss_g_all = self.loss_lamada[
'intentsity_loss'] * loss_int + self.loss_lamada[
'gradient_loss'] * loss_gd + self.loss_lamada[
'opticalflow_loss'] * loss_op + self.loss_lamada[
'gan_loss'] * loss_g_adv
self.optim_G.zero_grad()
loss_g_all.backward()
self.optim_G.step()
self.loss_meter_G.update(loss_g_all.detach())
if self.config.TRAIN.adversarial.scheduler.use:
self.lr_g.step()
#---------update optim_D ---------------
self.set_requires_grad(self.D, True)
self.optim_D.zero_grad()
real_d = self.D(torch.cat([target, flow_gt], dim=1))
fake_d = self.D(torch.cat([target, output_flow_G.detach()], dim=1))
loss_d_1 = self.gan_loss(real_d, True)
loss_d_2 = self.gan_loss(fake_d, False)
loss_d = (loss_d_1 + loss_d_2) * 0.5
loss_d.backward()
self.optim_D.step()
if self.config.TRAIN.adversarial.scheduler.use:
self.lr_d.step()
self.loss_meter_D.update(loss_d.detach())
# ======================End==================
self.batch_time.update(time.time() - start)
if (current_step % self.log_step == 0):
msg = 'Step: [{0}/{1}]\t' \
'Type: {cae_type}\t' \
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t' \
'Speed: {speed:.1f} samples/s\t' \
'Data: {data_time.val:.3f}s ({data_time.avg:.3f}s)\t' \
'Loss_G: {losses_G.val:.5f} ({losses_G.avg:.5f})\t' \
'Loss_D:{losses_D.val:.5f}({losses_D.avg:.5f})'.format(current_step, self.max_steps, cae_type=self.kwargs['model_type'], batch_time=self.batch_time, speed=self.config.TRAIN.batch_size/self.batch_time.val, data_time=self.data_time,losses_G=self.loss_meter_G, losses_D=self.loss_meter_D)
self.logger.info(msg)
writer.add_scalar('Train_loss_G', self.loss_meter_G.val, global_steps)
writer.add_scalar('Train_loss_D', self.loss_meter_D.val, global_steps)
if (current_step % self.vis_step == 0):
vis_objects = OrderedDict()
vis_objects['train_target_flow'] = flow_gt.detach()
vis_objects['train_output_flow_G'] = output_flow_G.detach()
vis_objects['train_target_frame'] = target.detach()
vis_objects['train_output_frame_G'] = output_frame_G.detach()
training_vis_images(vis_objects, writer, global_steps)
global_steps += 1
# reset start
start = time.time()
self.saved_model = {'G': self.G, 'D': self.D}
self.saved_optimizer = {
'optim_G': self.optim_G,
'optim_D': self.optim_D
}
self.saved_loss = {
'loss_G': self.loss_meter_G.val,
'loss_D': self.loss_meter_D.val
}
self.kwargs['writer_dict']['global_steps_{}'.format(
self.kwargs['model_type'])] = global_steps