def __init__(self, flownet_backbone): super(convAE, self).__init__() self.generator = UNet(input_channels=12, output_channel=3) self.discriminator = PixelDiscriminator(input_nc=3) self.flownet_backbone = flownet_backbone if flownet_backbone == '2sd': self.flow_net = FlowNet2SD() else: self.flow_net = lite_flow.Network()
def train(frame_num, layer_nums, input_channels, output_channels, discriminator_num_filters, bn=False, pretrain=False, generator_pretrain_path=None, discriminator_pretrain_path=None): generator = UNet(n_channels=input_channels, layer_nums=layer_nums, output_channel=output_channels, bn=bn) discriminator = PixelDiscriminator(output_channels, discriminator_num_filters, use_norm=False) generator = generator.cuda() discriminator = discriminator.cuda() flow_network = Network() flow_network.load_state_dict(torch.load(lite_flow_model_path)) flow_network.cuda().eval() adversarial_loss = Adversarial_Loss().cuda() discriminate_loss = Discriminate_Loss().cuda() gd_loss = Gradient_Loss(alpha, num_channels).cuda() op_loss = Flow_Loss().cuda() int_loss = Intensity_Loss(l_num).cuda() step = 0 if not pretrain: generator.apply(weights_init_normal) discriminator.apply(weights_init_normal) else: assert (generator_pretrain_path != None and discriminator_pretrain_path != None) generator.load_state_dict(torch.load(generator_pretrain_path)) discriminator.load_state_dict(torch.load(discriminator_pretrain_path)) step = int(generator_pretrain_path.split('-')[-1]) print('pretrained model loaded!') print('initializing the model with Generator-Unet {} layers,' 'PixelDiscriminator with filters {} '.format( layer_nums, discriminator_num_filters)) optimizer_G = torch.optim.Adam(generator.parameters(), lr=g_lr) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=d_lr) writer = SummaryWriter(writer_path) dataset = img_dataset.ano_pred_Dataset(training_data_folder, frame_num) dataset_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True) test_dataset = img_dataset.ano_pred_Dataset(testing_data_folder, frame_num) test_dataloader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True) for epoch in range(epochs): for (input, _), (test_input, _) in zip(dataset_loader, test_dataloader): # generator = generator.train() # discriminator = discriminator.train() target = input[:, -1, :, :, :].cuda() input = input[:, :-1, ] input_last = input[:, -1, ].cuda() input = input.view(input.shape[0], -1, input.shape[-2], input.shape[-1]).cuda() 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() #------- update optim_G -------------- G_output = generator(input) pred_flow_esti_tensor = torch.cat([input_last, G_output], 1) gt_flow_esti_tensor = torch.cat([input_last, target], 1) flow_gt = batch_estimate(gt_flow_esti_tensor, flow_network) flow_pred = batch_estimate(pred_flow_esti_tensor, flow_network) g_adv_loss = adversarial_loss(discriminator(G_output)) g_op_loss = op_loss(flow_pred, flow_gt) g_int_loss = int_loss(G_output, target) g_gd_loss = gd_loss(G_output, target) g_loss = lam_adv * g_adv_loss + lam_gd * g_gd_loss + lam_op * g_op_loss + lam_int * g_int_loss optimizer_G.zero_grad() g_loss.backward() optimizer_G.step() train_psnr = psnr_error(G_output, target) #----------- update optim_D ------- optimizer_D.zero_grad() d_loss = discriminate_loss(discriminator(target), discriminator(G_output.detach())) #d_loss.requires_grad=True d_loss.backward() optimizer_D.step() #----------- cal psnr -------------- test_generator = generator.eval() test_output = test_generator(test_input) test_psnr = psnr_error(test_output, test_target).cuda() if step % 10 == 0: print("[{}/{}]: g_loss: {} d_loss {}".format( step, epoch, g_loss, d_loss)) print('\t gd_loss {}, op_loss {}, int_loss {} ,'.format( g_gd_loss, g_op_loss, g_int_loss)) print('\t train psnr{},test_psnr {}'.format( train_psnr, test_psnr)) writer.add_scalar('psnr/train_psnr', train_psnr, global_step=step) writer.add_scalar('psnr/test_psnr', test_psnr, global_step=step) writer.add_scalar('total_loss/g_loss', g_loss, global_step=step) writer.add_scalar('total_loss/d_loss', d_loss, global_step=step) writer.add_scalar('g_loss/adv_loss', g_adv_loss, global_step=step) writer.add_scalar('g_loss/op_loss', g_op_loss, global_step=step) writer.add_scalar('g_loss/int_loss', g_int_loss, global_step=step) writer.add_scalar('g_loss/gd_loss', g_gd_loss, global_step=step) writer.add_image('image/train_target', target[0], global_step=step) writer.add_image('image/train_output', G_output[0], global_step=step) writer.add_image('image/test_target', test_target[0], global_step=step) writer.add_image('image/test_output', test_output[0], global_step=step) step += 1 if step % 500 == 0: utils.saver(generator.state_dict(), model_generator_save_path, step, max_to_save=10) utils.saver(discriminator.state_dict(), model_discriminator_save_path, step, max_to_save=10) if step >= 2000: print('==== begin evaluate the model of {} ===='.format( model_generator_save_path + '-' + str(step))) auc = evaluate(frame_num=5, layer_nums=4, input_channels=12, output_channels=3, model_path=model_generator_save_path + '-' + str(step), evaluate_name='compute_auc') writer.add_scalar('results/auc', auc, global_step=step)
def train(config): #### set the save and log path #### svname = args.name if svname is None: svname = config['train_dataset_type'] + '_' + config[ 'generator'] + '_' + config['flow_model'] if args.tag is not None: svname += '_' + args.tag save_path = os.path.join('./save', svname) utils.set_save_path(save_path) utils.set_log_path(save_path) writer = SummaryWriter(os.path.join(save_path, 'tensorboard')) yaml.dump(config, open(os.path.join(save_path, 'classifier_config.yaml'), 'w')) #### make datasets #### # train train_folder = config['dataset_path'] + config[ 'train_dataset_type'] + "/training/frames" test_folder = config['dataset_path'] + config[ 'train_dataset_type'] + "/testing/frames" # Loading dataset train_dataset_args = config['train_dataset_args'] test_dataset_args = config['test_dataset_args'] train_dataset = VadDataset(train_folder, transforms.Compose([ transforms.ToTensor(), ]), resize_height=train_dataset_args['h'], resize_width=train_dataset_args['w'], time_step=train_dataset_args['t_length'] - 1) test_dataset = VadDataset(test_folder, transforms.Compose([ transforms.ToTensor(), ]), resize_height=test_dataset_args['h'], resize_width=test_dataset_args['w'], time_step=test_dataset_args['t_length'] - 1) train_dataloader = DataLoader( train_dataset, batch_size=train_dataset_args['batch_size'], shuffle=True, num_workers=train_dataset_args['num_workers'], drop_last=True) test_dataloader = DataLoader(test_dataset, batch_size=test_dataset_args['batch_size'], shuffle=False, num_workers=test_dataset_args['num_workers'], drop_last=False) # for test---- prepare labels labels = np.load('./data/frame_labels_' + config['test_dataset_type'] + '.npy') if config['test_dataset_type'] == 'shanghai': labels = np.expand_dims(labels, 0) videos = OrderedDict() videos_list = sorted(glob.glob(os.path.join(test_folder, '*'))) labels_list = [] label_length = 0 psnr_list = {} for video in sorted(videos_list): # video_name = video.split('/')[-1] # windows video_name = os.path.split(video)[-1] videos[video_name] = {} videos[video_name]['path'] = video videos[video_name]['frame'] = glob.glob(os.path.join(video, '*.jpg')) videos[video_name]['frame'].sort() videos[video_name]['length'] = len(videos[video_name]['frame']) labels_list = np.append( labels_list, labels[0][4 + label_length:videos[video_name]['length'] + label_length]) label_length += videos[video_name]['length'] psnr_list[video_name] = [] # Model setting num_unet_layers = 4 discriminator_num_filters = [128, 256, 512, 512] # for gradient loss alpha = 1 # for int loss l_num = 2 pretrain = False if config['generator'] == 'cycle_generator_convlstm': ngf = 64 netG = 'resnet_6blocks' norm = 'instance' no_dropout = False init_type = 'normal' init_gain = 0.02 gpu_ids = [] generator = define_G(train_dataset_args['c'], train_dataset_args['c'], ngf, netG, norm, not no_dropout, init_type, init_gain, gpu_ids) elif config['generator'] == 'unet': # generator = UNet(n_channels=train_dataset_args['c']*(train_dataset_args['t_length']-1), # layer_nums=num_unet_layers, output_channel=train_dataset_args['c']) model = PreAE(train_dataset_args['c'], train_dataset_args['t_length'], **config['model_args']) else: raise Exception('The generator is not implemented') # generator = torch.load('save/avenue_cycle_generator_convlstm_flownet2_0103/generator-epoch-199.pth') discriminator = PixelDiscriminator(train_dataset_args['c'], discriminator_num_filters, use_norm=False) # discriminator = torch.load('save/avenue_cycle_generator_convlstm_flownet2_0103/discriminator-epoch-199.pth') # if not pretrain: # generator.apply(weights_init_normal) # discriminator.apply(weights_init_normal) # if use flownet # if config['flow_model'] == 'flownet2': # flownet2SD_model_path = 'flownet2/FlowNet2_checkpoint.pth.tar' # flow_network = FlowNet2(args).eval() # flow_network.load_state_dict(torch.load(flownet2SD_model_path)['state_dict']) # elif config['flow_model'] == 'liteflownet': # lite_flow_model_path = 'liteFlownet/network-sintel.pytorch' # flow_network = Network().eval() # flow_network.load_state_dict(torch.load(lite_flow_model_path)) # different range with the source version, should change lam_int = 1.0 * 2 lam_gd = 1.0 * 2 # here we use no flow loss lam_op = 0 # 2.0 lam_adv = 0.05 adversarial_loss = Adversarial_Loss() discriminate_loss = Discriminate_Loss() gd_loss = Gradient_Loss(alpha, train_dataset_args['c']) op_loss = Flow_Loss() int_loss = Intensity_Loss(l_num) step = 0 utils.log('initializing the model with Generator-Unet {} layers,' 'PixelDiscriminator with filters {} '.format( num_unet_layers, discriminator_num_filters)) g_lr = 0.0002 d_lr = 0.00002 optimizer_G = torch.optim.Adam(generator.parameters(), lr=g_lr) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=d_lr) # # optimizer setting # params_encoder = list(generator.encoder.parameters()) # params_decoder = list(generator.decoder.parameters()) # params = params_encoder + params_decoder # optimizer, lr_scheduler = utils.make_optimizer( # params, config['optimizer'], config['optimizer_args']) # # loss_func_mse = nn.MSELoss(reduction='none') # parallel if muti-gpus if torch.cuda.is_available(): generator.cuda() discriminator.cuda() # # if use flownet # flow_network.cuda() adversarial_loss.cuda() discriminate_loss.cuda() gd_loss.cuda() op_loss.cuda() int_loss.cuda() if config.get('_parallel'): generator = nn.DataParallel(generator) discriminator = nn.DataParallel(discriminator) # if use flownet # flow_network = nn.DataParallel(flow_network) adversarial_loss = nn.DataParallel(adversarial_loss) discriminate_loss = nn.DataParallel(discriminate_loss) gd_loss = nn.DataParallel(gd_loss) op_loss = nn.DataParallel(op_loss) int_loss = nn.DataParallel(int_loss) # Training utils.log('Start train') max_accuracy = 0 base_channel_num = train_dataset_args['c'] * ( train_dataset_args['t_length'] - 1) save_epoch = 5 if config['save_epoch'] is None else config['save_epoch'] for epoch in range(config['epochs']): generator.train() for j, imgs in enumerate( tqdm(train_dataloader, desc='train', leave=False)): imgs = imgs.cuda() input = imgs[:, :-1, ] input_last = input[:, -1, ] target = imgs[:, -1, ] # input = input.view(input.shape[0], -1, input.shape[-2],input.shape[-1]) # only for debug # input0=imgs[:, 0,] # input1 = imgs[:, 1, ] # gt_flow_esti_tensor = torch.cat([input0, input1], 1) # flow_gt = batch_estimate(gt_flow_esti_tensor, flow_network)[0] # objectOutput = open('./out_train.flo', 'wb') # np.array([80, 73, 69, 72], np.uint8).tofile(objectOutput) # np.array([flow_gt.size(2), flow_gt.size(1)], np.int32).tofile(objectOutput) # np.array(flow_gt.detach().cpu().numpy().transpose(1, 2, 0), np.float32).tofile(objectOutput) # objectOutput.close() # break # ------- update optim_G -------------- outputs = generator(input) # pred_flow_tensor = torch.cat([input_last, outputs], 1) # gt_flow_tensor = torch.cat([input_last, target], 1) # flow_pred = batch_estimate(pred_flow_tensor, flow_network) # flow_gt = batch_estimate(gt_flow_tensor, flow_network) # if you want to use flownet2SD, comment out the part in front # #### if use flownet #### # pred_flow_esti_tensor = torch.cat([input_last.view(-1,3,1,input.shape[-2],input.shape[-1]), # outputs.view(-1,3,1,input.shape[-2],input.shape[-1])], 2) # gt_flow_esti_tensor = torch.cat([input_last.view(-1,3,1,input.shape[-2],input.shape[-1]), # target.view(-1,3,1,input.shape[-2],input.shape[-1])], 2) # flow_gt=flow_network(gt_flow_esti_tensor*255.0) # flow_pred=flow_network(pred_flow_esti_tensor*255.0) ############################## # g_op_loss = op_loss(flow_pred, flow_gt) ## flow loss g_op_loss = 0 g_adv_loss = adversarial_loss(discriminator(outputs)) g_int_loss = int_loss(outputs, target) g_gd_loss = gd_loss(outputs, target) g_loss = lam_adv * g_adv_loss + lam_gd * g_gd_loss + lam_op * g_op_loss + lam_int * g_int_loss optimizer_G.zero_grad() g_loss.backward() optimizer_G.step() train_psnr = utils.psnr_error(outputs, target) # ----------- update optim_D ------- optimizer_D.zero_grad() d_loss = discriminate_loss(discriminator(target), discriminator(outputs.detach())) d_loss.backward() optimizer_D.step() # break # lr_scheduler.step() utils.log('----------------------------------------') utils.log('Epoch:' + str(epoch + 1)) utils.log('----------------------------------------') utils.log("g_loss: {} d_loss {}".format(g_loss, d_loss)) utils.log('\t gd_loss {}, op_loss {}, int_loss {} ,'.format( g_gd_loss, g_op_loss, g_int_loss)) utils.log('\t train psnr{}'.format(train_psnr)) # Testing utils.log('Evaluation of ' + config['test_dataset_type']) for video in sorted(videos_list): # video_name = video.split('/')[-1] video_name = os.path.split(video)[-1] psnr_list[video_name] = [] generator.eval() video_num = 0 # label_length += videos[videos_list[video_num].split('/')[-1]]['length'] label_length = videos[os.path.split( videos_list[video_num])[1]]['length'] for k, imgs in enumerate( tqdm(test_dataloader, desc='test', leave=False)): if k == label_length - 4 * (video_num + 1): video_num += 1 label_length += videos[os.path.split( videos_list[video_num])[1]]['length'] imgs = imgs.cuda() input = imgs[:, :-1, ] target = imgs[:, -1, ] # input = input.view(input.shape[0], -1, input.shape[-2], input.shape[-1]) outputs = generator(input) mse_imgs = int_loss((outputs + 1) / 2, (target + 1) / 2).item() # psnr_list[videos_list[video_num].split('/')[-1]].append(utils.psnr(mse_imgs)) psnr_list[os.path.split(videos_list[video_num])[1]].append( utils.psnr(mse_imgs)) # Measuring the abnormality score and the AUC anomaly_score_total_list = [] for video in sorted(videos_list): # video_name = video.split('/')[-1] video_name = os.path.split(video)[1] anomaly_score_total_list += utils.anomaly_score_list( psnr_list[video_name]) anomaly_score_total_list = np.asarray(anomaly_score_total_list) accuracy = utils.AUC(anomaly_score_total_list, np.expand_dims(1 - labels_list, 0)) utils.log('The result of ' + config['test_dataset_type']) utils.log('AUC: ' + str(accuracy * 100) + '%') # Save the model if epoch % save_epoch == 0 or epoch == config['epochs'] - 1: # torch.save(model, os.path.join( # save_path, 'model-epoch-{}.pth'.format(epoch))) torch.save( generator, os.path.join(save_path, 'generator-epoch-{}.pth'.format(epoch))) torch.save( discriminator, os.path.join(save_path, 'discriminator-epoch-{}.pth'.format(epoch))) if accuracy > max_accuracy: torch.save(generator, os.path.join(save_path, 'generator-max')) torch.save(discriminator, os.path.join(save_path, 'discriminator-max')) utils.log('----------------------------------------') utils.log('Training is finished')
def train(config): #### set the save and log path #### save_path = config['save_path'] utils.set_save_path(save_path) utils.set_log_path(save_path) writer = SummaryWriter(os.path.join(config['save_path'], 'tensorboard')) yaml.dump(config, open(os.path.join(config['save_path'], 'classifier_config.yaml'), 'w')) device = torch.device('cuda:' + args.gpu) #### make datasets #### # train train_folder = config['dataset_path'] + config['train_dataset_type'] + "/training/frames" test_folder = config['dataset_path'] + config['train_dataset_type'] + "/testing/frames" # Loading dataset train_dataset_args = config['train_dataset_args'] test_dataset_args = config['test_dataset_args'] train_dataset = VadDataset(args,video_folder= train_folder, bbox_folder = config['train_bboxes_path'], flow_folder=config['train_flow_path'], transform=transforms.Compose([transforms.ToTensor()]), resize_height=train_dataset_args['h'], resize_width=train_dataset_args['w'], dataset=config['train_dataset_type'], time_step=train_dataset_args['t_length'] - 1, device=device) test_dataset = VadDataset(args,video_folder= test_folder, bbox_folder = config['test_bboxes_path'], flow_folder=config['test_flow_path'], transform=transforms.Compose([transforms.ToTensor()]), resize_height=train_dataset_args['h'], resize_width=train_dataset_args['w'], dataset=config['train_dataset_type'], time_step=train_dataset_args['t_length'] - 1, device=device) train_dataloader = DataLoader(train_dataset, batch_size=train_dataset_args['batch_size'], shuffle=True, num_workers=train_dataset_args['num_workers'], drop_last=True) test_dataloader = DataLoader(test_dataset, batch_size=test_dataset_args['batch_size'], shuffle=False, num_workers=test_dataset_args['num_workers'], drop_last=False) # for test---- prepare labels labels = np.load('./data/frame_labels_' + config['test_dataset_type'] + '.npy') if config['test_dataset_type'] == 'shanghai': labels = np.expand_dims(labels, 0) videos = OrderedDict() videos_list = sorted(glob.glob(os.path.join(test_folder, '*'))) labels_list = [] label_length = 0 psnr_list = {} for video in sorted(videos_list): video_name = video.split('/')[-1] videos[video_name] = {} videos[video_name]['path'] = video videos[video_name]['frame'] = glob.glob(os.path.join(video, '*.jpg')) videos[video_name]['frame'].sort() videos[video_name]['length'] = len(videos[video_name]['frame']) labels_list = np.append(labels_list, labels[0][4 + label_length:videos[video_name]['length'] + label_length]) label_length += videos[video_name]['length'] psnr_list[video_name] = [] # Model setting num_unet_layers = 4 discriminator_num_filters = [128, 256, 512, 512] # for gradient loss alpha = 1 # for int loss l_num = 2 pretrain = False if config['generator'] == 'cycle_generator_convlstm': ngf = 64 netG = 'resnet_6blocks' norm = 'instance' no_dropout = False init_type = 'normal' init_gain = 0.02 gpu_ids = [] model = define_G(train_dataset_args['c'], train_dataset_args['c'], ngf, netG, norm, not no_dropout, init_type, init_gain, gpu_ids) elif config['generator'] == 'unet': # generator = UNet(n_channels=train_dataset_args['c']*(train_dataset_args['t_length']-1), # layer_nums=num_unet_layers, output_channel=train_dataset_args['c']) model = PreAE(train_dataset_args['c'], train_dataset_args['t_length'], **config['model_args']) else: raise Exception('The generator is not implemented') # generator = torch.load('save/avenue_cycle_generator_convlstm_flownet2_0103/generator-epoch-199.pth') if config['use_D']: discriminator=PixelDiscriminator(train_dataset_args['c'],discriminator_num_filters,use_norm=False) optimizer_D = torch.optim.Adam(discriminator.parameters(),lr=0.00002) # optimizer setting params_encoder = list(model.parameters()) params_decoder = list(model.parameters()) params = params_encoder + params_decoder optimizer_G, lr_scheduler = utils.make_optimizer( params, config['optimizer'], config['optimizer_args']) # set loss, different range with the source version, should change lam_int = 1.0 * 2 lam_gd = 1.0 * 2 # TODO here we use no flow loss # lam_op = 0 # 2.0 # op_loss = Flow_Loss() adversarial_loss = Adversarial_Loss() # TODO if use adv lam_adv = 0.05 discriminate_loss = Discriminate_Loss() alpha = 1 l_num = 2 gd_loss = Gradient_Loss(alpha, train_dataset_args['c']) int_loss = Intensity_Loss(l_num) object_loss = ObjectLoss(device, l_num) # parallel if muti-gpus if torch.cuda.is_available(): model.cuda() if config['use_D']: discriminator.cuda() if config.get('_parallel'): model = nn.DataParallel(model) if config['use_D']: discriminator = nn.DataParallel(discriminator) # Training utils.log('Start train') max_frame_AUC, max_roi_AUC = 0,0 base_channel_num = train_dataset_args['c'] * (train_dataset_args['t_length'] - 1) save_epoch = 5 if config['save_epoch'] is None else config['save_epoch'] for epoch in range(config['epochs']): model.train() for j, (imgs, bbox, flow) in enumerate(tqdm(train_dataloader, desc='train', leave=False)): imgs = imgs.cuda() flow = flow.cuda() # input = imgs[:, :-1, ].view(imgs.shape[0], -1, imgs.shape[-2], imgs.shape[-1]) input = imgs[:, :-1, ] target = imgs[:, -1, ] outputs = model(input) if config['use_D']: g_adv_loss = adversarial_loss(discriminator(outputs)) else: g_adv_loss = 0 g_object_loss = object_loss(outputs, target, flow, bbox) # g_int_loss = int_loss(outputs, target) g_gd_loss = gd_loss(outputs, target) g_loss = lam_adv * g_adv_loss + lam_gd * g_gd_loss + lam_int * g_object_loss optimizer_G.zero_grad() g_loss.backward() optimizer_G.step() train_psnr = utils.psnr_error(outputs,target) # ----------- update optim_D ------- if config['use_D']: optimizer_D.zero_grad() d_loss = discriminate_loss(discriminator(target), discriminator(outputs.detach())) d_loss.backward() optimizer_D.step() lr_scheduler.step() utils.log('----------------------------------------') utils.log('Epoch:' + str(epoch + 1)) utils.log('----------------------------------------') utils.log('Loss: Reconstruction {:.6f}'.format(g_loss.item())) # Testing utils.log('Evaluation of ' + config['test_dataset_type']) # Save the model if epoch % save_epoch == 0 or epoch == config['epochs'] - 1: if not os.path.exists(save_path): os.makedirs(save_path) if not os.path.exists(os.path.join(save_path, "models")): os.makedirs(os.path.join(save_path, "models")) # TODO frame_AUC = ObjectLoss_evaluate(test_dataloader, model, labels_list, videos, dataset=config['test_dataset_type'],device = device, frame_height = train_dataset_args['h'], frame_width=train_dataset_args['w'], is_visual=False, mask_labels_path = config['mask_labels_path'], save_path = os.path.join(save_path, "./final"), labels_dict=labels) torch.save(model.state_dict(), os.path.join(save_path, 'models/model-epoch-{}.pth'.format(epoch))) if config['use_D']: torch.save(discriminator.state_dict(), os.path.join(save_path, 'models/discrominator-epoch-{}.pth'.format(epoch))) else: frame_AUC = ObjectLoss_evaluate(test_dataloader, model, labels_list, videos, dataset=config['test_dataset_type'],device=device, frame_height = train_dataset_args['h'], frame_width=train_dataset_args['w']) utils.log('The result of ' + config['test_dataset_type']) utils.log("AUC: {}%".format(frame_AUC*100)) if frame_AUC > max_frame_AUC: max_frame_AUC = frame_AUC # TODO torch.save(model.state_dict(), os.path.join(save_path, 'models/max-frame_auc-model.pth')) if config['use_D']: torch.save(discriminator.state_dict(), os.path.join(save_path, 'models/discrominator-epoch-{}.pth'.format(epoch))) # evaluate(test_dataloader, model, labels_list, videos, int_loss, config['test_dataset_type'], test_bboxes=config['test_bboxes'], # frame_height = train_dataset_args['h'], frame_width=train_dataset_args['w'], # is_visual=True, mask_labels_path = config['mask_labels_path'], save_path = os.path.join(save_path, "./frame_best"), labels_dict=labels) utils.log('----------------------------------------') utils.log('Training is finished') utils.log('max_frame_AUC: {}'.format(max_frame_AUC))
parser.add_argument('--iters', default=40000, type=int, help='The total iteration number.') parser.add_argument('--resume', default=None, type=str, help='The pre-trained model to resume training with, pass \'latest\' or the model name.') parser.add_argument('--save_interval', default=1000, type=int, help='Save the model every [save_interval] iterations.') parser.add_argument('--val_interval', default=1000, type=int, help='Evaluate the model every [val_interval] iterations, pass -1 to disable.') parser.add_argument('--show_flow', default=False, action='store_true', help='If True, the first batch of ground truth optic flow could be visualized and saved.') parser.add_argument('--flownet', default='lite', type=str, help='lite: LiteFlownet, 2sd: FlowNet2SD.') args = parser.parse_args() train_cfg = update_config(args, mode='train') train_cfg.print_cfg() generator = UNet(input_channels=12, output_channel=3).cuda() discriminator = PixelDiscriminator(input_nc=3).cuda() optimizer_G = torch.optim.Adam(generator.parameters(), lr=train_cfg.g_lr) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=train_cfg.d_lr) if train_cfg.resume: generator.load_state_dict(torch.load(train_cfg.resume)['net_g']) discriminator.load_state_dict(torch.load(train_cfg.resume)['net_d']) optimizer_G.load_state_dict(torch.load(train_cfg.resume)['optimizer_g']) optimizer_D.load_state_dict(torch.load(train_cfg.resume)['optimizer_d']) print(f'Pre-trained generator and discriminator have been loaded.\n') else: generator.apply(weights_init_normal) discriminator.apply(weights_init_normal) print('Generator and discriminator are going to be trained from scratch.\n') assert train_cfg.flownet in ('lite', '2sd'), 'Flow net only supports LiteFlownet or FlowNet2SD currently.'