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