def custom_setup(self):
self.ovr_model = OneVsRestClassifier(LinearSVC(random_state = 0), n_jobs=16)
# basic meter
self.loss_meter_ABC = AverageMeter(name='loss_ABC')
self.ovr_model_path = os.path.join(self.config.TRAIN.model_output, f'ocae_cfg@{self.config_name}#{self.verbose}.npy')
self.cluster_dataset_keys = self.train_dataloaders_dict['cluster_dataset_dict'].keys()
def custom_setup(self):
# create loss meters
self.loss_meter_G = AverageMeter(name='Loss_G')
self.loss_meter_D = AverageMeter(name='Loss_D')
self.optical = ParamSet(
name='optical',
size=self.config.DATASET.optical_size,
output_format=self.config.DATASET.optical_format)
class ANOPREDTrainer(BaseTrainer):
NAME = ["ANOPRED.TRAIN"]
def custom_setup(self):
# basic meter
self.loss_meter_G = AverageMeter(name='Loss_G')
self.loss_meter_D = AverageMeter(name='Loss_D')
# others
self.optical = ParamSet(
name='optical',
size=self.config.DATASET.optical_size,
output_format=self.config.DATASET.optical_format)
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
class OCAETrainer(BaseTrainer):
NAME = ["OCAE.TRAIN"]
def custom_setup(self):
self.ovr_model = OneVsRestClassifier(LinearSVC(random_state = 0), n_jobs=16)
# basic meter
self.loss_meter_ABC = AverageMeter(name='loss_ABC')
self.ovr_model_path = os.path.join(self.config.TRAIN.model_output, f'ocae_cfg@{self.config_name}#{self.verbose}.npy')
self.cluster_dataset_keys = self.train_dataloaders_dict['cluster_dataset_dict'].keys()
# import ipdb; ipdb.set_trace()
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 custom_setup(self):
# basic meter
self.loss_predmeter_G = AverageMeter(name='loss_pred_G')
self.loss_predmeter_D = AverageMeter(name='loss_pred_D')
self.loss_refinemeter_G = AverageMeter(name='loss_refine_G')
self.loss_refinemeter_D = AverageMeter(name='loss_refine_D')
class ANOPCNTrainer(BaseTrainer):
NAME = ["ANOPCN.TRAIN"]
def custom_setup(self):
# basic meter
self.loss_predmeter_G = AverageMeter(name='loss_pred_G')
self.loss_predmeter_D = AverageMeter(name='loss_pred_D')
self.loss_refinemeter_G = AverageMeter(name='loss_refine_G')
self.loss_refinemeter_D = AverageMeter(name='loss_refine_D')
def train(self, current_step):
# Pytorch [N, C, D, H, W]
# initialize
self.set_requires_grad(self.F, False)
self.set_requires_grad(self.G, True)
self.set_requires_grad(self.D, True)
self.D.train()
self.G.train()
self.F.eval()
dynamic_steps = self.steps.param['dynamic_steps']
temp_step = current_step % dynamic_steps[2]
if temp_step in range(dynamic_steps[0], dynamic_steps[1]):
self.train_pcm(current_step)
elif temp_step in range(dynamic_steps[1], dynamic_steps[2]):
self.train_erm(current_step)
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_erm(self, current_step):
# Pytorch [N, C, D, H, W]
# initialize
start = time.time()
if self.kwargs['parallel']:
self.set_requires_grad(self.G.module.erm, True)
self.set_requires_grad(self.G.module.pcm, False)
else:
self.set_requires_grad(self.G.erm, True)
self.set_requires_grad(self.G.pcm, False)
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) # 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_refineframe_G = self.G(input_data, target)
gtFlowEstim = torch.cat([pred_last, target], 1).cuda()
predFlowEstim = torch.cat([pred_last, output_refineframe_G], 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_refineframe_G), True)
loss_op = self.OpticalflowSqrtLoss(predFlow, gtFlow)
loss_int = self.IntentsityLoss(output_refineframe_G, target)
loss_gd = self.GradientLoss(output_refineframe_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_refinemeter_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_refineframe_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.optim_D.step()
if self.config.TRAIN.adversarial.scheduler.use:
self.optimizer_D_scheduler.step()
self.loss_refinemeter_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_refinemeter_G, self.loss_refinemeter_D])
self.logger.info(msg)
writer.add_scalar('Train_loss_G', self.loss_refinemeter_G.val,
global_steps)
writer.add_scalar('Train_loss_D', self.loss_refinemeter_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_refineframe_G':
output_refineframe_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_refinemeter_G.val, 'loss_D':self.loss_refinemeter_D.val}
self.saved_loss['loss_G'] = self.loss_refinemeter_G.val
self.saved_loss['loss_D'] = self.loss_refinemeter_D.val
self.kwargs['writer_dict']['global_steps_{}'.format(
self.kwargs['model_type'])] = global_steps
def __init__(self, **kwargs):
"""Initialization Method.
Args:
kwargs(dict): the default will have:
model_dict: The directionary of the moddel
config:
parallel=parallel_flag,
pretrain=cfg.MODEL.PRETRAINED.USE
verbose=args.verbose,
time_stamp=time_stamp,
model_type=cfg.MODEL.NAME,
config_name=cfg_name,
hooks=hooks,
evaluate_function=evaluate_function,
final_output_dir=final_output_dir,
cpu=args.cpu
"""
self._hooks = []
self._register_hooks(kwargs['hooks'])
self.config = kwargs['config']
# devices
self.engine_gpus = self.config.SYSTEM.gpus
self.model = kwargs['model_dict']
if self.config.VAL.model_file == '':
raise Exception("Not have the Trained model file")
else:
self.model_path = self.config.VAL.model_file
# # get the optimizer
# self.optimizer = kwargs['optimizer_dict']
# # get the loss_fucntion
# self.loss_function = kwargs['loss_function_dict']
# basic meter, will be deprecated in the future.
self.batch_time = AverageMeter(name='batch_time')
self.data_time = AverageMeter(name='data_time')
# others
self.verbose = kwargs['verbose']
self.config_name = kwargs['config_name']
self.kwargs = kwargs
self.normalize = ParamSet(
name='normalize',
train={
'use': self.config.ARGUMENT.train.normal.use,
'mean': self.config.ARGUMENT.train.normal.mean,
'std': self.config.ARGUMENT.train.normal.std
},
val={
'use': self.config.ARGUMENT.val.normal.use,
'mean': self.config.ARGUMENT.val.normal.mean,
'std': self.config.ARGUMENT.val.normal.std
})
self.evaluate_function = kwargs['evaluate_function']
# hypyer-parameters of loss
# self.loss_lamada = kwargs['loss_lamada']
# the lr scheduler
# self.lr_scheduler_dict = kwargs['lr_scheduler_dict']
# initialize the saved objects
# None
# Get the models
for item_key in self.model.keys():
attr_name = str(item_key)
if self.kwargs['parallel']:
temp_model = self.data_parallel(self.model[item_key])
else:
temp_model = self.model[item_key].cuda()
self.__setattr__(attr_name, temp_model)
# get the losses
for item_key in self.loss_function.keys():
attr_name = str(item_key)
self.__setattr__(attr_name, self.loss_function[attr_name])
self.custom_setup()
self.load_model(self.model_path)
def __init__(self, **kwargs):
"""Initialization Method.
Args:
defaults(tuple): the default will have:
0 0->model:{'Generator':net_g, 'Driscriminator':net_d, 'FlowNet':net_flow}
- 1->train_dataloader: the dataloader # Will be deprecated in the future
- 2->val_dataloader: the dataloader # Will be deprecated in the future
1 -->dataloader_dict: the dict of all the dataloader will be used in the process
2 3->optimizer:{'optimizer_g':op_g, 'optimizer_d'}
3 4->loss_function: {'g_adverserial_loss':.., 'd_adverserial_loss':..., 'gradient_loss':.., 'opticalflow_loss':.., 'intentsity_loss':.. }
4 5->logger: the logger of the whole training process
5 6->config: the config object of the whole process
kwargs(dict): the default will have:
model_dict: The directionary of the moddel
dataloaders_dict:
optimizer_dict:
loss_function_dict:
logger:
cfg:
parallel=parallel_flag,
pretrain=cfg.MODEL.PRETRAINED.USE
verbose=args.verbose,
time_stamp=time_stamp,
model_type=cfg.MODEL.NAME,
writer_dict=writer_dict,
config_name=cfg_name,
loss_lamada=loss_lamada,
hooks=hooks,
evaluate_function=evaluate_function,
lr_scheduler_dict=lr_scheduler_dict,
final_output_dir=final_output_dir,
cpu=args.cpu
"""
self._hooks = []
# self._eval_hooks = []
self._register_hooks(kwargs['hooks'])
# logger & config
# self.logger = defaults[4]
self.config = kwargs['config']
# devices
self.engine_gpus = self.config.SYSTEM.gpus
self.model = kwargs['model_dict']
if kwargs['pretrain']:
self.load_pretrain()
dataloaders_dict = kwargs['dataloaders_dict']
self.val_dataloaders_dict = dataloaders_dict['train']
self.val_dataset_keys = list(
dataloaders_dict['test']['general_dataset_dict'].keys())
# get the optimizer
self.optimizer = kwargs['optimizer_dict']
# get the loss_fucntion
self.loss_function = kwargs['loss_function_dict']
# basic meter
self.batch_time = AverageMeter(name='batch_time')
self.data_time = AverageMeter(name='data_time')
# others
self.verbose = kwargs['verbose']
self.accuarcy = 0.0 # to store the accuracy varies from epoch to epoch
self.config_name = kwargs['config_name']
self.result_path = ''
self.kwargs = kwargs
self.normalize = ParamSet(
name='normalize',
train={
'use': self.config.ARGUMENT.train.normal.use,
'mean': self.config.ARGUMENT.train.normal.mean,
'std': self.config.ARGUMENT.train.normal.std
},
val={
'use': self.config.ARGUMENT.val.normal.use,
'mean': self.config.ARGUMENT.val.normal.mean,
'std': self.config.ARGUMENT.val.normal.std
})
self.steps = ParamSet(name='steps',
log=self.config.VAL.log_step,
vis=self.config.VAL.vis_step,
eval=self.config.TRAIN.eval_step,
max=self.config.TRAIN.max_steps)
self.evaluate_function = kwargs['evaluate_function']
# hypyer-parameters of loss
self.loss_lamada = kwargs['loss_lamada']
# the lr scheduler
self.lr_scheduler_dict = kwargs['lr_scheduler_dict']
# initialize the saved objects
# None
# Get the models
for item_key in self.model.keys():
attr_name = str(item_key)
if self.kwargs['parallel']:
temp_model = self.data_parallel(self.model[item_key])
else:
temp_model = self.model[item_key].cuda()
self.__setattr__(attr_name, temp_model)
# get the optimizer
# None
# get the losses
for item_key in self.loss_function.keys():
attr_name = str(item_key)
self.__setattr__(attr_name, self.loss_function[attr_name])
self.custom_setup()
class MATrainer(BaseTrainer):
"""
G
D_frame
D_pattern
AE_act
AE_obj
PatternNet
"""
NAME = ["MA.TRAIN"]
def custom_setup(self):
# create loss meters
self.loss_meter_G = AverageMeter(name='Loss_G')
self.loss_meter_D = AverageMeter(name='Loss_D')
self.optical = ParamSet(
name='optical',
size=self.config.DATASET.optical_size,
output_format=self.config.DATASET.optical_format)
# import ipdb; ipdb.set_trace()
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 __init__(self, *defaults, **kwargs):
"""
Args:
defaults(tuple): the default will have:
0 0->model:{'Generator':net_g, 'Driscriminator':net_d, 'FlowNet':net_flow}
- 1->train_dataloader: the dataloader # Will be deprecated in the future
- 2->val_dataloader: the dataloader # Will be deprecated in the future
1 -->dataloader_dict: the dict of all the dataloader will be used in the process
2 3->optimizer:{'optimizer_g':op_g, 'optimizer_d'}
3 4->loss_function: {'g_adverserial_loss':.., 'd_adverserial_loss':..., 'gradient_loss':.., 'opticalflow_loss':.., 'intentsity_loss':.. }
4 5->logger: the logger of the whole training process
5 6->config: the config object of the whole process
kwargs(dict): the default will have:
verbose(str):
parallel(bool): True-> data parallel
pertrain(bool): True-> use the pretarin model
dataloaders_dict: will to replace the train_dataloader and test_dataloader
extra param:
test_dataset_keys: the dataset keys of each video
test_dataset_dict: the dataset dict of whole test videos
"""
self._hooks = []
# self._eval_hooks = []
self._register_hooks(kwargs['hooks'])
# logger & config
self.logger = defaults[4]
self.config = defaults[5]
self.model = defaults[0]
if kwargs['pretrain']:
self.load_pretrain()
dataloaders_dict = defaults[1]
self.val_dataloaders_dict = dataloaders_dict['train']
self.val_dataset_keys = list(
dataloaders_dict['test']['general_dataset_dict'].keys())
# get the optimizer
self.optimizer = defaults[2]
# get the loss_fucntion
self.loss_function = defaults[3]
# basic meter
self.batch_time = AverageMeter(name='batch_time')
self.data_time = AverageMeter(name='data_time')
# others
self.verbose = kwargs['verbose']
self.accuarcy = 0.0 # to store the accuracy varies from epoch to epoch
self.config_name = kwargs['config_name']
self.result_path = ''
self.kwargs = kwargs
self.normalize = ParamSet(
name='normalize',
train={
'use': self.config.ARGUMENT.train.normal.use,
'mean': self.config.ARGUMENT.train.normal.mean,
'std': self.config.ARGUMENT.train.normal.std
},
val={
'use': self.config.ARGUMENT.val.normal.use,
'mean': self.config.ARGUMENT.val.normal.mean,
'std': self.config.ARGUMENT.val.normal.std
})
self.steps = ParamSet(name='steps',
log=self.config.VAL.log_step,
vis=self.config.VAL.vis_step,
eval=self.config.TRAIN.eval_step,
max=self.config.TRAIN.max_steps)
self.evaluate_function = kwargs['evaluate_function']
# hypyer-parameters of loss
self.loss_lamada = kwargs['loss_lamada']
# the lr scheduler
self.lr_scheduler_dict = kwargs['lr_scheduler_dict']
# initialize the saved objects
# None
# Get the models
for item_key in self.model.keys():
attr_name = str(item_key)
if self.kwargs['parallel']:
temp_model = self.data_parallel(self.model[item_key])
else:
temp_model = self.model[item_key].cuda()
self.__setattr__(attr_name, temp_model)
# get the optimizer
# None
# get the losses
for item_key in self.loss_function.keys():
attr_name = str(item_key)
self.__setattr__(attr_name, self.loss_function[attr_name])
self.custom_setup()
def custom_setup(self):
# basic meter
self.loss_meter_STAE = AverageMeter(name='loss_STAE')
class STAETrainer(BaseTrainer):
NAME = ["STAE.TRAIN"]
def custom_setup(self):
# basic meter
self.loss_meter_STAE = AverageMeter(name='loss_STAE')
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
