def train(self, current_step):
# Pytorch [N, C, D, H, W]
# initialize
start = time.time()
self.set_requires_grad(self.F, False)
self.set_requires_grad(self.G, True)
self.set_requires_grad(self.D, True)
self.G.train()
self.D.train()
self.F.eval()
writer = self.kwargs['writer_dict']['writer']
global_steps = self.kwargs['writer_dict']['global_steps_{}'.format(
self.kwargs['model_type'])]
# get the data
data, anno, meta = 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.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_pred_G = self.G(input_data)
# import ipdb; ipdb.set_trace()
predFlowEstim = torch.cat([input_last, output_pred_G], 1)
gtFlowEstim = torch.cat([input_last, target], 1)
gtFlow_vis, gtFlow = flow_batch_estimate(
self.F,
gtFlowEstim,
self.normalize.param['train'],
output_format=self.config.DATASET.optical_format,
optical_size=self.config.DATASET.optical_size)
predFlow_vis, predFlow = flow_batch_estimate(
self.F,
predFlowEstim,
self.normalize.param['train'],
output_format=self.config.DATASET.optical_format,
optical_size=self.config.DATASET.optical_size)
loss_g_adv = self.GANLoss(self.D(output_pred_G), True)
loss_op = self.OpticalflowLoss(predFlow, gtFlow)
loss_int = self.IntentsityLoss(output_pred_G, target)
loss_gd = self.GradientLoss(output_pred_G, target)
loss_g_all = self.loss_lamada[
'IntentsityLoss'] * loss_int + self.loss_lamada[
'GradientLoss'] * loss_gd + self.loss_lamada[
'OpticalflowLoss'] * loss_op + self.loss_lamada[
'GANLoss'] * loss_g_adv
self.optimizer_G.zero_grad()
loss_g_all.backward()
self.optimizer_G.step()
# record
self.loss_meter_G.update(loss_g_all.detach())
if self.config.TRAIN.adversarial.scheduler.use:
self.optimizer_G_scheduler.step()
#---------update optim_D ---------------
self.set_requires_grad(self.D, True)
self.optimizer_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.GANLoss(temp_t, True)
loss_d_2 = self.GANLoss(temp_g, False)
loss_d = (loss_d_1 + loss_d_2) * 0.5
# loss_d.sum().backward()
loss_d.backward()
self.optimizer_D.step()
if self.config.TRAIN.adversarial.scheduler.use:
self.optimizer_D_scheduler.step()
self.loss_meter_D.update(loss_d.detach())
# ======================End==================
self.batch_time.update(time.time() - start)
if (current_step % self.steps.param['log'] == 0):
msg = make_info_message(current_step, self.steps.param['max'],
self.kwargs['model_type'], self.batch_time,
self.config.TRAIN.batch_size,
self.data_time,
[self.loss_meter_G, 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.steps.param['vis'] == 0):
vis_objects = OrderedDict({
'train_target':
target.detach(),
'train_output_pred_G':
output_pred_G.detach(),
'train_gtFlow':
gtFlow_vis.detach(),
'train_predFlow':
predFlow_vis.detach()
})
tensorboard_vis_images(vis_objects, writer, global_steps,
self.normalize.param['train'])
global_steps += 1
# reset start
start = time.time()
# self.saved_model = {'G':self.G, 'D':self.D}
self.saved_model['G'] = self.G
self.saved_model['D'] = self.D
# self.saved_optimizer = {'optim_G': self.optimizer_G, 'optim_D': self.optimizer_D}
self.saved_optimizer['optimizer_G'] = self.optimizer_G
self.saved_optimizer['optimizer_D'] = self.optimizer_D
# self.saved_loss = {'loss_G':self.loss_meter_G.val, 'loss_D':self.loss_meter_D.val}
self.saved_loss['loss_G'] = self.loss_meter_G.val
self.saved_loss['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
'''
self.engine.set_requires_grad(self.engine.MemAE, False)
self.engine.MemAE.eval()
tb_writer = self.engine.kwargs['writer_dict']['writer']
global_steps = self.engine.kwargs['writer_dict'][
'global_steps_{}'.format(self.engine.kwargs['model_type'])]
frame_num = self.engine.config.DATASET.val.clip_length
clip_step = self.engine.config.DATASET.val.clip_step
psnr_records = []
score_records = []
# total = 0
num_videos = 0
random_video_sn = torch.randint(0, len(self.engine.test_dataset_keys),
(1, ))
# calc the score for the test dataset
for sn, video_name in enumerate(self.engine.test_dataset_keys):
num_videos += 1
# need to improve
dataset = self.engine.test_dataset_dict[video_name]
len_dataset = dataset.pics_len
test_iters = len_dataset - frame_num + 1
# test_iters = len_dataset // clip_step
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))
# scores = np.empty(shape=(len_dataset,),dtype=np.float32)
scores = torch.zeros(len_dataset)
# scores = [0.0 for i in range(len_dataset)]
for clip_sn, (test_input, anno, meta) in enumerate(data_loader):
test_target = test_input.cuda()
time_len = test_input.shape[2]
output, _ = self.engine.MemAE(test_target)
clip_score = reconstruction_loss(output, test_target)
# scores[test_counter*time_len:(test_counter + 1)*time_len] = clip_score.squeeze(0)
if (frame_num + test_counter) > len_dataset:
temp = test_counter + frame_num - len_dataset
scores[test_counter:len_dataset] = clip_score[temp:]
else:
scores[test_counter:(frame_num +
test_counter)] = clip_score
test_counter += 1
if sn == random_video_sn and (clip_sn in vis_range):
vis_objects = OrderedDict({
'memae_eval_clip':
test_target.detach(),
'memae_eval_clip_hat':
output.detach()
})
tensorboard_vis_images(
vis_objects,
tb_writer,
global_steps,
normalize=self.engine.normalize.param['val'])
if test_counter >= test_iters:
# 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)
smin = min(scores)
# normal_scores = np.array([(1.0 - np.divide(s-smin, smax)) for s in scores])
normal_scores = (1.0 - torch.div(
scores - smin, smax - smin)).detach().cpu().numpy()
normal_scores = np.clip(normal_scores, 0, None)
score_records.append(normal_scores)
print(f'finish test video set {video_name}')
break
self.engine.pkl_path = save_score_results(
self.engine.config,
self.engine.logger,
verbose=self.engine.verbose,
config_name=self.engine.config_name,
current_step=current_step,
time_stamp=self.engine.kwargs["time_stamp"],
score=score_records)
results = self.engine.evaluate_function(
self.engine.pkl_path, self.engine.logger, self.engine.config,
self.engine.config.DATASET.score_type)
self.engine.logger.info(results)
tb_writer.add_text('amc: AUC of ROC curve', f'auc is {results.auc}',
global_steps)
return results.auc
def train_pcm(self, current_step):
# Pytorch [N, C, D, H, W]
# initialize
start = time.time()
if self.kwargs['parallel']:
self.set_requires_grad(self.G.module.pcm, True)
self.set_requires_grad(self.G.module.erm, False)
else:
self.set_requires_grad(self.G.pcm, True)
self.set_requires_grad(self.G.erm, False)
writer = self.kwargs['writer_dict']['writer']
global_steps = self.kwargs['writer_dict']['global_steps_{}'.format(
self.kwargs['model_type'])]
# get the data
data, anno, meta = next(
self._train_loader_iter) # the core for dataloader
self.data_time.update(time.time() - start)
# base on the D to get each frame
target = data[:, :, -1, :, :].cuda() # t frame
pred_last = data[:, :, -2, :, :].cuda() # t-1 frame
input_data = data[:, :, :-1, :, :].cuda() # 0 ~ t-1 frame
# input_data = data.cuda() # 0 ~ t frame
# True Process =================Start===================
#---------update optim_G ---------
self.set_requires_grad(self.D, False)
output_predframe_G, _ = self.G(input_data, target)
predFlowEstim = torch.cat([pred_last, output_predframe_G], 1).cuda()
gtFlowEstim = torch.cat([pred_last, target], 1).cuda()
gtFlow_vis, gtFlow = flow_batch_estimate(
self.F,
gtFlowEstim,
self.normalize.param['train'],
output_format=self.config.DATASET.optical_format,
optical_size=self.config.DATASET.optical_size)
predFlow_vis, predFlow = flow_batch_estimate(
self.F,
predFlowEstim,
self.normalize.param['train'],
output_format=self.config.DATASET.optical_format,
optical_size=self.config.DATASET.optical_size)
loss_g_adv = self.GANLoss(self.D(output_predframe_G), True)
loss_op = self.OpticalflowSqrtLoss(predFlow, gtFlow)
loss_int = self.IntentsityLoss(output_predframe_G, target)
loss_gd = self.GradientLoss(output_predframe_G, target)
loss_g_all = self.loss_lamada['IntentsityLoss'] * loss_int + self.loss_lamada['GradientLoss'] * loss_gd + \
self.loss_lamada['OpticalflowSqrtLoss'] * loss_op + self.loss_lamada['GANLoss'] * loss_g_adv
self.optimizer_G.zero_grad()
loss_g_all.backward()
self.optimizer_G.step()
# record
self.loss_predmeter_G.update(loss_g_all.detach())
if self.config.TRAIN.adversarial.scheduler.use:
self.optimizer_G_scheduler.step()
#---------update optim_D ---------------
self.set_requires_grad(self.D, True)
self.optimizer_D.zero_grad()
temp_t = self.D(target)
temp_g = self.D(output_predframe_G.detach())
loss_d_1 = self.GANLoss(temp_t, True)
loss_d_2 = self.GANLoss(temp_g, False)
loss_d = (loss_d_1 + loss_d_2) * 0.5
loss_d.backward()
self.optimizer_D.step()
if self.config.TRAIN.adversarial.scheduler.use:
self.optimizer_D_scheduler.step()
self.loss_predmeter_D.update(loss_d.detach())
# ======================End==================
self.batch_time.update(time.time() - start)
if (current_step % self.steps.param['log'] == 0):
msg = make_info_message(
current_step, self.steps.param['max'],
self.kwargs['model_type'], self.batch_time,
self.config.TRAIN.batch_size, self.data_time,
[self.loss_predmeter_G, self.loss_predmeter_D])
self.logger.info(msg)
writer.add_scalar('Train_loss_G', self.loss_predmeter_G.val,
global_steps)
writer.add_scalar('Train_loss_D', self.loss_predmeter_D.val,
global_steps)
if (current_step % self.steps.param['vis'] == 0):
vis_objects = OrderedDict({
'train_target_flow':
gtFlow_vis.detach(),
'train_pred_flow':
predFlow_vis.detach(),
'train_target_frame':
target.detach(),
'train_output_predframe_G':
output_predframe_G.detach()
})
tensorboard_vis_images(vis_objects, writer, global_steps,
self.normalize.param['train'])
global_steps += 1
# reset start
start = time.time()
self.saved_model = {'G': self.G, 'D': self.D}
self.saved_optimizer = {
'optim_G': self.optimizer_G,
'optim_D': self.optimizer_D
}
self.saved_loss = {
'loss_G': self.loss_predmeter_G.val,
'loss_D': self.loss_predmeter_D.val
}
self.kwargs['writer_dict']['global_steps_{}'.format(
self.kwargs['model_type'])] = global_steps
def train(self,current_step):
# Pytorch [N, C, D, H, W]
# initialize
start = time.time()
self.set_requires_grad(self.A, True)
self.set_requires_grad(self.B, True)
self.set_requires_grad(self.C, True)
self.set_requires_grad(self.Detector, False)
self.A.train()
self.B.train()
self.C.train()
self.Detector.eval()
writer = self.kwargs['writer_dict']['writer']
global_steps = self.kwargs['writer_dict']['global_steps_{}'.format(self.kwargs['model_type'])]
# get the data
data, anno, meta = next(self._train_loader_iter) # the core for dataloader
self.data_time.update(time.time() - start)
# base on the D to get each frame
# in this method, D = 3 and not change
future = data[:, :, -1, :, :].cuda() # t+1 frame
current = data[:, :, 1, :, :].cuda() # t frame
past = data[:, :, 0, :, :].cuda() # t-1 frame
bboxs = get_batch_dets(self.Detector, current)
# this method is based on the objects to train the model insted of frames
for index, bbox in enumerate(bboxs):
if bbox.numel() == 0:
bbox = bbox.new_zeros([1, 4])
# get the crop objects
input_currentObject_B, _ = multi_obj_grid_crop(current[index], bbox)
future_object, _ = multi_obj_grid_crop(future[index], bbox)
# future2current = torch.stack([future_object, input_currentObject_B], dim=1)
current2future = torch.stack([input_currentObject_B, future_object], dim=1)
past_object, _ = multi_obj_grid_crop(past[index], bbox)
# current2past = torch.stack([input_currentObject_B, past_object], dim=1)
past2current = torch.stack([past_object, input_currentObject_B], dim=1)
_, _, input_objectGradient_A = frame_gradient(current2future)
input_objectGradient_A = input_objectGradient_A.sum(1)
_, _, input_objectGradient_C = frame_gradient(past2current)
input_objectGradient_C = input_objectGradient_C.sum(1)
# import ipdb; ipdb.set_trace()
# True Process =================Start===================
# original_A = (0.3 * input_objectGradient_A[:,0] + 0.59 * input_objectGradient_A[:,1] + 0.11 * input_objectGradient_A[:,2]).unsqueeze(1)
# original_B = (0.3 * input_currentObject_B[:,0] + 0.59 * input_currentObject_B[:,1] + 0.11 * input_currentObject_B[:,2]).unsqueeze(1)
# original_C = (0.3 * input_objectGradient_C[:,0] + 0.59 * input_objectGradient_C[:,1] + 0.11 * input_objectGradient_C[:,2]).unsqueeze(1)
_, output_recGradient_A, original_A = self.A(input_objectGradient_A)
_, output_recObject_B, original_B = self.B(input_currentObject_B)
_, output_recGradient_C, original_C = self.C(input_objectGradient_C)
# import ipdb; ipdb.set_trace()
# loss_A = self.a_loss(output_recGradient_A, input_objectGradient_A)
# loss_B = self.b_loss(output_recObject_B, input_currentObject_B)
# loss_C = self.c_loss(output_recGradient_C, input_objectGradient_C)
loss_A = self.ALoss(output_recGradient_A, original_A)
loss_B = self.BLoss(output_recObject_B, original_B)
loss_C = self.CLoss(output_recGradient_C, original_C)
loss_all = self.loss_lamada['ALoss'] * loss_A + self.loss_lamada['BLoss'] * loss_B + self.loss_lamada['CLoss'] * loss_C
self.optimizer_ABC.zero_grad()
loss_all.backward()
self.optimizer_ABC.step()
# record
self.loss_meter_ABC.update(loss_all.detach())
if self.config.TRAIN.general.scheduler.use:
self.optimizer_ABC_scheduler.step()
# ======================End==================
self.batch_time.update(time.time() - start)
if (current_step % self.steps.param['log'] == 0):
msg = make_info_message(current_step, self.steps.param['max'], self.kwargs['model_type'], self.batch_time, self.config.TRAIN.batch_size, self.data_time, [self.loss_meter_ABC])
self.logger.info(msg)
writer.add_scalar('Train_loss_ABC', self.loss_meter_ABC.val, global_steps)
if (current_step % self.steps.param['vis'] == 0):
vis_objects = OrderedDict({
'train_input_objectGradient_A': input_objectGradient_A.detach(),
'train_input_currentObject_B': input_currentObject_B.detach(),
'train_input_objectGradient_C': input_objectGradient_C.detach(),
'train_output_recGradient_A': output_recGradient_A.detach(),
'train_output_recObject_B': output_recObject_B.detach(),
'train_output_recGradient_C': output_recGradient_C.detach()
})
tensorboard_vis_images(vis_objects, writer, global_steps, self.normalize.param['train'])
global_steps += 1
# reset start
start = time.time()
# self.saved_model = {'A':self.A, 'B':self.B, 'C':self.C}
# self.saved_model['A'] = self.A
# self.saved_model['B'] = self.B
# self.saved_model['C'] = self.C
# # self.saved_optimizer = {'optim_ABC': self.optim_ABC}
# self.saved_optimizer['optimizer_ABC'] = self.optimizer_ABC
# # self.saved_loss = {'loss_ABC':self.loss_meter_ABC.val}
# self.saved_loss['loss_ABC'] = self.loss_meter_ABC.val
self.saved_stuff['step'] = global_steps
self.saved_stuff['loss'] = self.loss_meter_ABC.val
self.saved_stuff['A'] = self.A
self.saved_stuff['B'] = self.B
self.saved_stuff['C'] = self.C
self.saved_stuff['optimizer_ABC'] = self.optimizer_ABC
self.kwargs['writer_dict']['global_steps_{}'.format(self.kwargs['model_type'])] = global_steps
def train(self, current_step):
# Pytorch [N, C, D, H, W]
# initialize
start = time.time()
self.set_requires_grad(self.F, False)
self.set_requires_grad(self.D, True)
self.set_requires_grad(self.G, True)
self.G.train()
self.D.train()
self.F.eval()
writer = self.kwargs['writer_dict']['writer']
global_steps = self.kwargs['writer_dict']['global_steps_{}'.format(
self.kwargs['model_type'])]
# get the data
data, anno, meta = 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
input_data = data[:, :, 0, :, :].cuda() # input(1-st) frame
target = data[:, :, 1, :, :].cuda() # target(2-nd) frame
# 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_vis, flow_gt = flow_batch_estimate(
self.F,
gt_flow_esti_tensor,
self.normalize.param['train'],
optical_size=self.config.DATASET.optical_size,
output_format=self.config.DATASET.optical_format)
fake_g = self.D(torch.cat([target, output_flow_G], dim=1))
loss_g_adv = self.GANLoss(fake_g, True)
loss_op = self.OpticalflowSqrtLoss(output_flow_G, flow_gt)
loss_int = self.IntentsityLoss(output_frame_G, target)
loss_gd = self.GradientLoss(output_frame_G, target)
loss_g_all = self.loss_lamada[
'IntentsityLoss'] * loss_int + self.loss_lamada[
'GradientLoss'] * loss_gd + self.loss_lamada[
'OpticalflowSqrtLoss'] * loss_op + self.loss_lamada[
'GANLoss'] * loss_g_adv
self.optimizer_G.zero_grad()
loss_g_all.backward()
self.optimizer_G.step()
self.loss_meter_G.update(loss_g_all.detach())
if self.config.TRAIN.adversarial.scheduler.use:
self.optimizer_G_scheduler.step()
#---------update optim_D ---------------
self.set_requires_grad(self.D, True)
self.optimizer_D.zero_grad()
# import ipdb; ipdb.set_trace()
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.GANLoss(real_d, True)
loss_d_2 = self.GANLoss(fake_d, False)
loss_d = (loss_d_1 + loss_d_2) * 0.5
loss_d.backward()
self.optimizer_D.step()
if self.config.TRAIN.adversarial.scheduler.use:
self.optimizer_D_scheduler.step()
self.loss_meter_D.update(loss_d.detach())
# ======================End==================
self.batch_time.update(time.time() - start)
if (current_step % self.steps.param['log'] == 0):
msg = make_info_message(current_step, self.steps.param['max'],
self.kwargs['model_type'], self.batch_time,
self.config.TRAIN.batch_size,
self.data_time,
[self.loss_meter_G, self.loss_meter_D])
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.steps.param['vis'] == 0):
temp = vis_optical_flow(
output_flow_G.detach(),
output_format=self.config.DATASET.optical_format,
output_size=(output_flow_G.shape[-2], output_flow_G.shape[-1]),
normalize=self.normalize.param['train'])
vis_objects = OrderedDict({
'train_target_flow':
flow_gt_vis.detach(),
'train_output_flow_G':
temp,
'train_target_frame':
target.detach(),
'train_output_frame_G':
output_frame_G.detach(),
})
tensorboard_vis_images(vis_objects, writer, global_steps,
self.normalize.param['train'])
global_steps += 1
# reset start
start = time.time()
# self.saved_model = {'G':self.G, 'D':self.D}
self.saved_model['G'] = self.G
self.saved_model['D'] = self.D
# self.saved_optimizer = {'optim_G': self.optimizer_G, 'optim_D': self.optimizer_D}
self.saved_optimizer['optimizer_G'] = self.optimizer_G
self.saved_optimizer['optimizer_D'] = self.optimizer_D
# self.saved_loss = {'loss_G':self.loss_meter_G.val, 'loss_D':self.loss_meter_D.val}
self.saved_loss['loss_G'] = self.loss_meter_G.val
self.saved_loss['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
'''
# self.trainer.set_requires_grad(self.trainer.A, False)
# self.trainer.set_requires_grad(self.trainer.B, False)
# self.trainer.set_requires_grad(self.trainer.C, False)
# self.trainer.set_requires_grad(self.trainer.Detector, False)
# self.trainer.A.eval()
# self.trainer.B.eval()
# self.trainer.C.eval()
# self.trainer.Detector.eval()
self.engine.set_all(False)
frame_num = self.engine.config.DATASET.test_clip_length
tb_writer = self.engine.kwargs['writer_dict']['writer']
global_steps = self.engine.kwargs['writer_dict']['global_steps_{}'.format(self.engine.kwargs['model_type'])]
score_records = []
# psnr_records = []
total = 0
random_video_sn = torch.randint(0, len(self.engine.test_dataset_keys), (1,))
# random_video_sn = 0
for sn, video_name in enumerate(self.engine.test_dataset_keys):
# _temp_test_folder = os.path.join(self.testing_data_folder, dir)
# need to improve
# dataset = AvenueTestOld(_temp_test_folder, clip_length=frame_num)
dataset = self.engine.test_dataset_dict[video_name]
len_dataset = dataset.pics_len
test_iters = len_dataset - frame_num + 1
test_counter = 0
# feature_record = []
data_loader = DataLoader(dataset=dataset, batch_size=1, shuffle=False, num_workers=1)
# import ipdb; ipdb.set_trace()
scores = np.empty(shape=(len_dataset,),dtype=np.float32)
# for test_input, _ in data_loader:
random_frame_sn = torch.randint(0, len_dataset,(1,))
for frame_sn, (test_input, anno, meta) in enumerate(data_loader):
feature_record_object = []
future = test_input[:, :, 2, :, :].cuda()
current = test_input[:, :, 1, :, :].cuda()
past = test_input[:, :, 0, :, :].cuda()
bboxs = get_batch_dets(self.engine.Detector, current)
for index, bbox in enumerate(bboxs):
# import ipdb; ipdb.set_trace()
if bbox.numel() == 0:
bbox = bbox.new_zeros([1,4])
# print('NO objects')
# continue
current_object, _ = multi_obj_grid_crop(current[index], bbox)
future_object, _ = multi_obj_grid_crop(future[index], bbox)
future2current = torch.stack([future_object, current_object], dim=1)
past_object, _ = multi_obj_grid_crop(past[index], bbox)
current2past = torch.stack([current_object, past_object], dim=1)
_, _, A_input = frame_gradient(future2current)
A_input = A_input.sum(1)
_, _, C_input = frame_gradient(current2past)
C_input = C_input.sum(1)
A_feature, temp_a, _ = self.engine.A(A_input)
B_feature, temp_b, _ = self.engine.B(current_object)
C_feature, temp_c, _ = self.engine.C(C_input)
# import ipdb; ipdb.set_trace()
if sn == random_video_sn and frame_sn == random_frame_sn:
vis_objects = OrderedDict({
'eval_oc_input_a': A_input.detach(),
'eval_oc_output_a': temp_a.detach(),
'eval_oc_input_b': current_object.detach(),
'eval_oc_output_b': temp_b.detach(),
'eval_oc_input_c': C_input.detach(),
'eval_oc_output_c': temp_c.detach(),
})
tensorboard_vis_images(vis_objects, tb_writer, global_steps, normalize=self.engine.normalize.param['val'])
A_flatten_feature = A_feature.flatten(start_dim=1)
B_flatten_feature = B_feature.flatten(start_dim=1)
C_flatten_feature = C_feature.flatten(start_dim=1)
ABC_feature = torch.cat([A_flatten_feature, B_flatten_feature, C_flatten_feature], dim=1).detach()
ABC_feature_s = torch.chunk(ABC_feature, ABC_feature.size(0), dim=0)
for abc_f in ABC_feature_s:
temp = abc_f.squeeze(0).cpu().numpy()
feature_record_object.append(temp)
predict_input = np.array(feature_record_object)
self.engine.ovr_model = joblib.load(self.engine.ovr_model_path)
g_i = self.engine.ovr_model.decision_function(predict_input) # the svm score of each object in one frame
frame_score = oc_score(g_i)
# test_psnr = psnr_error(g_output, test_target)
# test_psnr = test_psnr.tolist()
scores[test_counter+frame_num-1] = frame_score
test_counter += 1
total+=1
if test_counter >= test_iters:
scores[:frame_num-1]=scores[frame_num-1]
score_records.append(scores)
print(f'finish test video set {video_name}')
break
self.engine.pkl_path = save_score_results(self.engine.config, self.engine.logger, verbose=self.engine.verbose, config_name=self.engine.config_name, current_step=current_step, time_stamp=self.engine.kwargs["time_stamp"],score=score_records)
results = self.engine.evaluate_function(self.engine.pkl_path, self.engine.logger, self.engine.config)
self.engine.logger.info(results)
tb_writer.add_text('AUC of ROC curve', f'AUC is {results.auc:.5f}',global_steps)
return results.auc
def evaluate(self, current_step):
"""STAE evaluation method.
Evaluate the model base on some methods.
Args:
current_step: The current step at present
Returns:
results: The magnitude of the method based on this evaluation metric
"""
# Set basic things
self.engine.set_all(False) # eval mode
tb_writer = self.engine.kwargs['writer_dict']['writer']
global_steps = self.engine.kwargs['writer_dict'][
'global_steps_{}'.format(self.engine.kwargs['model_type'])]
frame_num = self.engine.config.DATASET.val.sampled_clip_length
score_records = []
# num_videos = 0
random_video_sn = torch.randint(0, len(self.engine.val_dataset_keys),
(1, ))
# calc the score for the test dataset
for sn, video_name in enumerate(self.engine.val_dataset_keys):
# num_videos += 1
# need to improve
dataloader = self.engine.val_dataloaders_dict[
'general_dataset_dict'][video_name]
len_dataset = dataloader.dataset.pics_len
test_iters = len_dataset - frame_num + 1
# test_iters = len_dataset // clip_step
test_counter = 0
vis_range = range(int(len_dataset * 0.5),
int(len_dataset * 0.5 + 5))
scores = np.empty(shape=(len_dataset, ), dtype=np.float32)
for clip_sn, (test_input, anno, meta) in enumerate(dataloader):
test_input = test_input.cuda()
# test_target = data[:,:,16:,:,:].cuda()
time_len = test_input.shape[2]
output, _ = self.engine.STAE(test_input)
clip_score = reconstruction_loss(output, test_input)
clip_score = clip_score.tolist()
if (frame_num + test_counter) > len_dataset:
temp = test_counter + frame_num - len_dataset
scores[test_counter:len_dataset] = clip_score[temp:]
else:
scores[test_counter:(frame_num +
test_counter)] = clip_score
test_counter += 1
if sn == random_video_sn and (clip_sn in vis_range):
vis_objects = OrderedDict({
'stae_eval_clip':
test_input.detach(),
'stae_eval_clip_hat':
output.detach()
})
tensorboard_vis_images(
vis_objects,
tb_writer,
global_steps,
normalize=self.engine.normalize.param['val'])
if test_counter >= test_iters:
# scores[:frame_num-1]=(scores[frame_num-1],) # fix the bug: TypeError: can only assign an iterable
smax = max(scores)
smin = min(scores)
normal_scores = np.array([(1.0 - np.divide(s - smin, smax))
for s in scores])
normal_scores = np.clip(normal_scores, 0, None)
score_records.append(normal_scores)
logger.info(f'Finish testing the video:{video_name}')
break
# Compute the metrics based on the model's results
self.engine.pkl_path = save_score_results(
score_records,
self.engine.config,
self.engine.logger,
verbose=self.engine.verbose,
config_name=self.engine.config_name,
current_step=current_step,
time_stamp=self.engine.kwargs["time_stamp"])
results = self.engine.evaluate_function.compute(
{'val': self.engine.pkl_path})
self.engine.logger.info(results)
# Write the metric into the tensorboard
tb_writer.add_text(
f'{self.engine.config.MODEL.name}: AUC of ROC curve',
f'auc is {results.avg_value}', global_steps)
return results.avg_value
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
'''
# self.trainer.set_requires_grad(self.trainer.F, False)
# self.trainer.set_requires_grad(self.trainer.G, False)
# self.trainer.set_requires_grad(self.trainer.D, False)
# self.trainer.G.eval()
# self.trainer.D.eval()
# self.trainer.F.eval()
self.engine.set_all(False) # eval mode
tb_writer = self.engine.kwargs['writer_dict']['writer']
global_steps = self.engine.kwargs['writer_dict'][
'global_steps_{}'.format(self.engine.kwargs['model_type'])]
frame_num = self.engine.config.DATASET.val.clip_length
psnr_records = []
score_records = []
# total = 0
# for dirs in video_dirs:
random_video_sn = torch.randint(0, len(self.engine.val_dataset_keys),
(1, ))
for sn, video_name in enumerate(self.engine.val_dataset_keys):
# need to improve
dataset = self.engine.val_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)
psnrs = np.empty(shape=(len_dataset, ), dtype=np.float32)
scores = np.empty(shape=(len_dataset, ), dtype=np.float32)
vis_range = range(int(len_dataset * 0.5),
int(len_dataset * 0.5 + 5))
for frame_sn, (test_input, anno, meta) in enumerate(data_loader):
test_target = test_input[:, :, -1, :, :].cuda()
test_input = test_input[:, :, :-1, :, :].reshape(
test_input.shape[0], -1, test_input.shape[-2],
test_input.shape[-1]).cuda()
g_output = self.engine.G(test_input)
test_psnr = psnr_error(g_output.detach(),
test_target,
hat=True)
test_psnr = test_psnr.tolist()
psnrs[test_counter + frame_num - 1] = test_psnr
scores[test_counter + frame_num - 1] = test_psnr
test_counter += 1
# total+=1
if sn == random_video_sn and (frame_sn in vis_range):
vis_objects = OrderedDict({
'anopred_eval_frame':
test_target.detach(),
'anopred_eval_frame_hat':
g_output.detach()
})
tensorboard_vis_images(
vis_objects,
tb_writer,
global_steps,
normalize=self.engine.normalize.param['val'])
if test_counter >= test_iters:
psnrs[:frame_num - 1] = psnrs[frame_num - 1]
scores[:frame_num - 1] = (scores[frame_num - 1], )
smax = max(scores)
smin = min(scores)
normal_scores = np.array(
[np.divide(s - smin, smax - smin) for s in scores])
normal_scores = np.clip(normal_scores, 0, None)
psnr_records.append(psnrs)
score_records.append(normal_scores)
# print(f'finish test video set {video_name}')
break
self.engine.pkl_path = save_score_results(
self.engine.config,
self.engine.logger,
verbose=self.engine.verbose,
config_name=self.engine.config_name,
current_step=current_step,
time_stamp=self.engine.kwargs["time_stamp"],
score=score_records,
psnr=psnr_records)
results = self.engine.evaluate_function(
self.engine.pkl_path, self.engine.logger, self.engine.config,
self.engine.config.DATASET.score_type)
self.engine.logger.info(results)
tb_writer.add_text('anopcn: AUC of ROC curve', f'auc is {results.auc}',
global_steps)
return results.auc
def evaluate(self, current_step):
"""AnoPCN evaluation method.
Evaluate the model base on some methods.
Args:
current_step: The current step at present
Returns:
results: The magnitude of the method based on this evaluation metric
"""
# Set basic things
self.engine.set_all(False)
tb_writer = self.engine.kwargs['writer_dict']['writer']
global_steps = self.engine.kwargs['writer_dict'][
'global_steps_{}'.format(self.engine.kwargs['model_type'])]
frame_num = self.engine.config.DATASET.test_clip_length
# psnr_records=[]
score_records = []
total = 0
# for dirs in video_dirs:
random_video_sn = torch.randint(0, len(self.engine.test_dataset_keys),
(1, ))
for sn, video_name in enumerate(self.engine.test_dataset_keys):
# need to improve
# dataset = self.engine.test_dataset_dict[video_name]
dataloader = self.engine.val_dataloaders_dict[
'general_dataset_dict'][video_name]
len_dataset = dataloader.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)
# import ipdb; ipdb.set_trace()
# psnrs = np.empty(shape=(len_dataset,),dtype=np.float32)
scores = np.empty(shape=(len_dataset, ), dtype=np.float32)
vis_range = range(int(len_dataset * 0.5),
int(len_dataset * 0.5 + 5))
for frame_sn, (test_input, anno, meta) in enumerate(dataloader):
test_target = test_input[:, :, -1, :, :].cuda()
test_input = test_input[:, :, :-1, :, :].cuda()
_, g_output = self.engine.G(test_input, test_target)
test_psnr = psnr_error(g_output, test_target, hat=False)
test_psnr = test_psnr.tolist()
# psnrs[test_counter+frame_num-1]=test_psnr
scores[test_counter + frame_num - 1] = test_psnr
if sn == random_video_sn and (frame_sn in vis_range):
vis_objects = OrderedDict({
'anopcn_eval_frame':
test_target.detach(),
'anopcn_eval_frame_hat':
g_output.detach()
})
# vis_objects['anopcn_eval_frame'] = test_target.detach()
# vis_objects['anopcn_eval_frame_hat'] = g_output.detach()
tensorboard_vis_images(
vis_objects,
tb_writer,
global_steps,
normalize=self.engine.normalize.param['val'])
test_counter += 1
total += 1
if test_counter >= test_iters:
# psnrs[:frame_num-1]=psnrs[frame_num-1]
scores[:frame_num - 1] = (scores[frame_num - 1], )
smax = max(scores)
smin = min(scores)
normal_scores = np.array(
[np.divide(s - smin, smax - smin) for s in scores])
normal_scores = np.clip(normal_scores, 0, None)
# psnr_records.append(psnrs)
score_records.append(normal_scores)
logger.info(f'finish test video set {video_name}')
break
# Compute the metrics based on the model's results
self.engine.pkl_path = save_score_results(
score_records,
self.engine.config,
self.engine.logger,
verbose=self.engine.verbose,
config_name=self.engine.config_name,
current_step=current_step,
time_stamp=self.engine.kwargs["time_stamp"])
results = self.engine.evaluate_function.compute(
{'val': self.engine.pkl_path})
self.engine.logger.info(results)
# Write the metric into the tensorboard
tb_writer.add_text(
f'{self.engine.config.MODEL.name}: AUC of ROC curve',
f'auc is {results.avg_value}', global_steps)
return results.auc
def train(self, current_step):
# Pytorch [N, C, D, H, W]
# initialize
start = time.time()
self.set_requires_grad(self.STAE, True)
self.STAE.train()
writer = self.kwargs['writer_dict']['writer']
global_steps = self.kwargs['writer_dict']['global_steps_{}'.format(
self.kwargs['model_type'])]
# get the data
data, anno, meta = next(
self._train_loader_iter) # the core for dataloader
self.data_time.update(time.time() - start)
# get the reconstruction and prediction video clip
time_len = data.shape[2]
rec_time = time_len // 2
input_rec = data[:, :, 0:rec_time, :, :].cuda() # 0 ~ t//2 frame
input_pred = data[:, :,
rec_time:time_len, :, :].cuda() # t//2 ~ t frame
# True Process =================Start===================
output_rec, output_pred = self.STAE(input_rec)
loss_rec = self.RecLoss(output_rec, input_rec)
loss_pred = self.WeightedPredLoss(output_pred, input_pred)
# print(f'loss_rec:{loss_rec}')
# print(f'loss_pred:{loss_pred}')
# loss_stae_all = self.loss_lamada['rec_loss'] * loss_rec + self.loss_lamada['weighted_pred_loss'] * loss_pred
loss_stae_all = self.loss_lamada[
'RecLoss'] * loss_rec + self.loss_lamada[
'WeightedPredLoss'] * loss_pred
# self.optim_STAE.zero_grad()
self.optimizer_STAE.zero_grad()
loss_stae_all.backward()
# self.optim_STAE.step()
self.optimizer_STAE.step()
self.loss_meter_STAE.update(loss_stae_all.detach())
if self.config.TRAIN.general.scheduler.use:
# self.lr_stae.step()
self.optimizer_STAE_scheduler.step()
# ======================End==================
self.batch_time.update(time.time() - start)
if (current_step % self.steps.param['log'] == 0):
msg = make_info_message(current_step, self.steps.param['max'],
self.kwargs['model_type'], self.batch_time,
self.config.TRAIN.batch_size,
self.data_time, [self.loss_meter_STAE])
self.logger.info(msg)
writer.add_scalar('Train_loss_STAE', self.loss_meter_STAE.val,
global_steps)
if (current_step % self.steps.param['vis'] == 0):
vis_objects = OrderedDict({
'train_output_rec': output_rec.detach(),
'train_output_pred': output_pred.detach(),
'train_input_rec': input_rec.detach(),
'train_input_pred': input_pred.detach()
})
tensorboard_vis_images(vis_objects, writer, global_steps,
self.normalize.param['train'])
global_steps += 1
# reset start
start = time.time()
# self.saved_model = {'STAE':self.STAE}
self.saved_model['STAE'] = self.STAE
# self.saved_optimizer = {'optim_STAE': self.optim_STAE}
# self.saved_optimizer = {'optim_STAE': self.optimizer_STAE}
self.saved_optimizer['optimizer_STAE'] = self.optimizer_STAE
# self.saved_loss = {'loss_STAE':self.loss_meter_STAE}
self.saved_loss['loss_STAE'] = self.loss_meter_STAE
self.kwargs['writer_dict']['global_steps_{}'.format(
self.kwargs['model_type'])] = global_steps
def evaluate(self, current_step):
"""AMC evaluation method.
Evaluate the model base on some methods.
Args:
current_step: The current step at present
Returns:
results: The magnitude of the method based on this evaluation metric
"""
# Set basic things
self.engine.set_all(False)
tb_writer = self.engine.kwargs['writer_dict']['writer']
global_steps = self.engine.kwargs['writer_dict'][
'global_steps_{}'.format(self.engine.kwargs['model_type'])]
frame_num = self.engine.config.DATASET.val.sampled_clip_length
# psnr_records=[]
score_records = []
# score_records_w=[]
w_dict = OrderedDict()
# total = 0
# calc the weight for the training set
w_video_dict = self.engine.train_dataloaders_dict['w_dataset_dict']
w_video_names = self.engine.train_dataloaders_dict[
'w_dataset_dict'].keys()
for video_name in w_video_names:
# dataset = self.engine.test_dataset_dict_w[video_name]
data_loader = w_video_dict[video_name]
len_dataset = data_loader.dataset.pics_len
test_iters = len_dataset - frame_num + 1
test_counter = 0
scores = [0.0 for i in range(len_dataset)]
for data, _, _ in data_loader:
input_data_test = data[:, :, 0, :, :].cuda()
target_test = data[:, :, 1, :, :].cuda()
# import ipdb; ipdb.set_trace()
output_flow_G, output_frame_G = self.engine.G(input_data_test)
gtFlowEstim = torch.cat([input_data_test, target_test], 1)
gtFlow_vis, gtFlow = flow_batch_estimate(
self.engine.F,
gtFlowEstim,
self.engine.normalize.param['val'],
output_format=self.engine.config.DATASET.optical_format,
optical_size=self.engine.config.DATASET.optical_size)
# import ipdb; ipdb.set_trace()
diff_appe, diff_flow = simple_diff(target_test, output_frame_G,
gtFlow, output_flow_G)
# patch_score_appe, patch_score_flow, _, _ = find_max_patch(diff_appe, diff_flow)
patch_score_appe, patch_score_flow = find_max_patch(
diff_appe, diff_flow)
scores[test_counter + frame_num -
1] = [patch_score_appe, patch_score_flow]
test_counter += 1
# print(test_counter)
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]
logger.info(
f'Finish calc the scores of training set {video_name} in step:{current_step}'
)
break
# import ipdb; ipdb.set_trace()
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)
logger.info(f'wf:{wf}, wi:{wi}')
# calc the score for the test dataset
num_videos = 0
random_video_sn = torch.randint(0, len(self.engine.test_dataset_keys),
(1, ))
for sn, video_name in enumerate(self.engine.test_dataset_keys):
num_videos += 1
# need to improve
dataloader = self.engine.val_dataloaders_dict[
'general_dataset_dict'][video_name]
len_dataset = dataloader.dataset.pics_len
test_iters = len_dataset - frame_num + 1
test_counter = 0
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)
for frame_sn, (data, anno, meta) in enumerate(dataloader):
test_input = data[:, :, 0, :, :].cuda()
test_target = data[:, :, 1, :, :].cuda()
g_output_flow, g_output_frame = self.engine.G(test_input)
gt_flow_esti_tensor = torch.cat([test_input, test_target], 1)
flow_gt_vis, flow_gt = flow_batch_estimate(
self.engine.F,
gt_flow_esti_tensor,
self.engine.param['val'],
output_format=self.engine.config.DATASET.optical_format,
optical_size=self.engine.config.DATASET.optical_size)
# test_psnr = psnr_error(g_output_frame, test_target)
score, _, _ = amc_score(test_target, g_output_frame, flow_gt,
g_output_flow, wf, wi)
# 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):
temp = vis_optical_flow(
g_output_flow.detach(),
output_format=self.engine.config.DATASET.
optical_format,
output_size=(g_output_flow.shape[-2],
g_output_flow.shape[-1]),
normalize=self.engine.normalize.param['val'])
vis_objects = OrderedDict({
'amc_eval_frame':
test_target.detach(),
'amc_eval_frame_hat':
g_output_frame.detach(),
'amc_eval_flow':
flow_gt_vis.detach(),
'amc_eval_flow_hat':
temp
})
tensorboard_vis_images(
vis_objects,
tb_writer,
global_steps,
normalize=self.engine.normalize.param['val'])
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])
normal_scores = np.clip(normal_scores, 0, None)
# psnr_records.append(psnrs)
score_records.append(normal_scores)
logger.info(f'Finish test video set {video_name}')
break
# Compute the metrics based on the model's results
self.engine.pkl_path = save_score_results(
score_records,
self.engine.config,
self.engine.logger,
verbose=self.engine.verbose,
config_name=self.engine.config_name,
current_step=current_step,
time_stamp=self.engine.kwargs["time_stamp"])
results = self.engine.evaluate_function.compute(
{'val': self.engine.pkl_path})
self.engine.logger.info(results)
# Write the metric into the tensorboard
tb_writer.add_text(
f'{self.engine.config.MODEL.name}: AUC of ROC curve',
f'auc is {results.auc}', global_steps)
return results.auc
