def train(args): start_t = time.time() params = get_train_options() params["exp_name"] = args.exp_name params["patch_num_point"] = 1024 params["batch_size"] = args.batch_size params['use_gan'] = args.use_gan if args.debug: params["nepoch"] = 2 params["model_save_interval"] = 3 params['model_vis_interval'] = 3 log_dir = os.path.join(params["model_save_dir"], args.exp_name) if os.path.exists(log_dir) == False: os.makedirs(log_dir) tb_logger = Logger(log_dir) trainloader = PUNET_Dataset(h5_file_path=params["dataset_dir"], split_dir=params['train_split']) #print(params["dataset_dir"]) num_workers = 4 train_data_loader = data.DataLoader(dataset=trainloader, batch_size=params["batch_size"], shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') G_model = Generator(params) G_model.apply(xavier_init) G_model = torch.nn.DataParallel(G_model).to(device) D_model = Discriminator(params, in_channels=3) D_model.apply(xavier_init) D_model = torch.nn.DataParallel(D_model).to(device) G_model.train() D_model.train() optimizer_D = Adam(D_model.parameters(), lr=params["lr_D"], betas=(0.9, 0.999)) optimizer_G = Adam(G_model.parameters(), lr=params["lr_G"], betas=(0.9, 0.999)) D_scheduler = MultiStepLR(optimizer_D, [50, 80], gamma=0.2) G_scheduler = MultiStepLR(optimizer_G, [50, 80], gamma=0.2) Loss_fn = Loss() print("preparation time is %fs" % (time.time() - start_t)) iter = 0 for e in range(params["nepoch"]): D_scheduler.step() G_scheduler.step() for batch_id, (input_data, gt_data, radius_data) in enumerate(train_data_loader): optimizer_G.zero_grad() optimizer_D.zero_grad() input_data = input_data[:, :, 0:3].permute(0, 2, 1).float().cuda() gt_data = gt_data[:, :, 0:3].permute(0, 2, 1).float().cuda() start_t_batch = time.time() output_point_cloud = G_model(input_data) repulsion_loss = Loss_fn.get_repulsion_loss( output_point_cloud.permute(0, 2, 1)) uniform_loss = Loss_fn.get_uniform_loss( output_point_cloud.permute(0, 2, 1)) #print(output_point_cloud.shape,gt_data.shape) emd_loss = Loss_fn.get_emd_loss( output_point_cloud.permute(0, 2, 1), gt_data.permute(0, 2, 1)) if params['use_gan'] == True: fake_pred = D_model(output_point_cloud.detach()) d_loss_fake = Loss_fn.get_discriminator_loss_single( fake_pred, label=False) d_loss_fake.backward() optimizer_D.step() real_pred = D_model(gt_data.detach()) d_loss_real = Loss_fn.get_discriminator_loss_single(real_pred, label=True) d_loss_real.backward() optimizer_D.step() d_loss = d_loss_real + d_loss_fake fake_pred = D_model(output_point_cloud) g_loss = Loss_fn.get_generator_loss(fake_pred) #print(repulsion_loss,uniform_loss,emd_loss) total_G_loss=params['uniform_w']*uniform_loss+params['emd_w']*emd_loss+ \ repulsion_loss*params['repulsion_w']+ g_loss*params['gan_w'] else: #total_G_loss = params['uniform_w'] * uniform_loss + params['emd_w'] * emd_loss + \ # repulsion_loss * params['repulsion_w'] total_G_loss=params['emd_w'] * emd_loss + \ repulsion_loss * params['repulsion_w'] #total_G_loss=emd_loss total_G_loss.backward() optimizer_G.step() current_lr_D = optimizer_D.state_dict()['param_groups'][0]['lr'] current_lr_G = optimizer_G.state_dict()['param_groups'][0]['lr'] tb_logger.scalar_summary('repulsion_loss', repulsion_loss.item(), iter) tb_logger.scalar_summary('uniform_loss', uniform_loss.item(), iter) tb_logger.scalar_summary('emd_loss', emd_loss.item(), iter) if params['use_gan'] == True: tb_logger.scalar_summary('d_loss', d_loss.item(), iter) tb_logger.scalar_summary('g_loss', g_loss.item(), iter) tb_logger.scalar_summary('lr_D', current_lr_D, iter) tb_logger.scalar_summary('lr_G', current_lr_G, iter) msg = "{:0>8},{}:{}, [{}/{}], {}: {},{}:{}".format( str(datetime.timedelta(seconds=round(time.time() - start_t))), "epoch", e, batch_id + 1, len(train_data_loader), "total_G_loss", total_G_loss.item(), "iter time", (time.time() - start_t_batch)) print(msg) if iter % params['model_vis_interval'] == 0 and iter > 0: np_pcd = output_point_cloud.permute( 0, 2, 1)[0].detach().cpu().numpy() #print(np_pcd.shape) img = (np.array(visualize_point_cloud(np_pcd)) * 255).astype( np.uint8) tb_logger.image_summary("images", img[np.newaxis, :], iter) gt_pcd = gt_data.permute(0, 2, 1)[0].detach().cpu().numpy() #print(gt_pcd.shape) gt_img = (np.array(visualize_point_cloud(gt_pcd)) * 255).astype(np.uint8) tb_logger.image_summary("gt", gt_img[np.newaxis, :], iter) input_pcd = input_data.permute(0, 2, 1)[0].detach().cpu().numpy() input_img = (np.array(visualize_point_cloud(input_pcd)) * 255).astype(np.uint8) tb_logger.image_summary("input", input_img[np.newaxis, :], iter) iter += 1 if (e + 1) % params['model_save_interval'] == 0 and e > 0: model_save_dir = os.path.join(params['model_save_dir'], params['exp_name']) if os.path.exists(model_save_dir) == False: os.makedirs(model_save_dir) D_ckpt_model_filename = "D_iter_%d.pth" % (e) G_ckpt_model_filename = "G_iter_%d.pth" % (e) D_model_save_path = os.path.join(model_save_dir, D_ckpt_model_filename) G_model_save_path = os.path.join(model_save_dir, G_ckpt_model_filename) torch.save(D_model.module.state_dict(), D_model_save_path) torch.save(G_model.module.state_dict(), G_model_save_path)
''' mtcnn+asf 准确率测评 ''' resize_width = 160 resize_height = 160 frame_interval = 3 # Number of frames after which to run face detection fps_display_interval = 5 # seconds frame_rate = 0 frame_count = 0 normalization = False # 是否标准化,默认否 data_path = "E:/testData/lfw" log = Logger('./logs/mtcnn_asf_eval.log', level='info') # 保存测试集人物信息:姓名,特征大小,特征,创建时间 class userInfo(): def __init__(self, username, feature_size, feature, create_time): self.username = username self.feature_size = feature_size self.feature = feature self.create_time = create_time if __name__ == '__main__': # 1.mtcnn初始化 face_detect = face_recognition.FaceDetection() # 初始化mtcnn
from utils.ParseExcel import ParseExcel from ConfigRead import * from action.PageAction import PageAction from selenium.common.exceptions import * # 导入所有异常类 from utils.Logger import Logger from utils.ParseYaml import ParseYaml from utils.WriterFile import YamlWrite from testcase.HTMLTestRunner_cn import HTMLTestRunner from datetime import datetime from openpyxl.styles import Font from BeautifulReport import BeautifulReport as bf, BeautifulReport import traceback import time import re import unittest logger = Logger('logger').getlog() class RunnerTestCase(unittest.TestCase): def setUp(self): self.parseyaml = ParseYaml() self.testdata_path = self.parseyaml.ReadParameter('ImportAddress') self.parseexcel = ParseExcel(self.testdata_path) self.pageaction = PageAction() self.sheetnames = self.parseexcel.wb.sheetnames self.parameter = CONFIG_PATH + 'Parameter.yaml' self.CaseNum = 0 # 创建六个字典,分别储存步骤测试结果,用例测试结果,用例测试时间,错误信息,截图信息,步骤测试时间 self.time_dic = {} self.result_dic = {}
num = 200 OUTLOG = 'statistics/log_gDQN' Base = '/home/ayal/Documents/gym/Code/' DQN_home = 'DQN_hockey/hockey_DDQN_deepmind/hockey_DQN_'+str(num)+'_V' DDPG_home = 'DDPG2/results' gDQN_home = 'DQN_hockey/hockey_numeric_3points/hockey_multinet1_decay_rate1.2_2_V' gDQNd = [Base+gDQN_home+'7/logs',Base+gDQN_home+'6/logs',Base+gDQN_home+'5/logs',Base+gDQN_home+'2/logs']#,Base+gDQN_home+'4/logs'] DDPGd = [Base+DDPG_home+'/logs',Base+DDPG_home+'1/logs',Base+DDPG_home+'2/logs',Base+DDPG_home+'3/logs',Base+DDPG_home+'4/logs'] DQNd = [Base+DQN_home+'1/logs',Base+DQN_home+'2/logs',Base+DQN_home+'3/logs',Base+DQN_home+'4/logs',Base+DQN_home+'5/logs'] numlogs = len(gDQNd) LEN = 100 L = Logger() L.AddNewLog('network_left') L.AddNewLog('network_left_up') L.AddNewLog('network_left_down') L.AddNewLog('network_middle') L.AddNewLog('network_middle_up') L.AddNewLog('network_middle_down') L.AddNewLog('network_right') L.AddNewLog('network_right_up') L.AddNewLog('network_right_down') L.AddNewLog('network_random') L.AddNewLog('network_random_up') L.AddNewLog('network_random_down') L.AddNewLog('estimated_value') L.AddNewLog('estimated_value_up') L.AddNewLog('estimated_value_down')
model_conf = Params(os.path.join(conf.conf_dir, conf.model.lower() + '.json')) model_conf.update_dict('exp_conf', conf.__dict__) np.random.seed(conf.seed) torch.random.manual_seed(conf.seed) device = torch.device('cuda') if torch.cuda.is_available() else torch.device( 'cpu') dataset = Dataset(data_dir=conf.data_dir, data_name=model_conf.data_name, train_ratio=model_conf.train_ratio, device=device) log_dir = os.path.join('saves', conf.model) logger = Logger(log_dir) model_conf.save(os.path.join(logger.log_dir, 'config.json')) eval_pos, eval_target = dataset.eval_data() item_popularity = dataset.item_popularity evaluator = Evaluator(eval_pos, eval_target, item_popularity, model_conf.top_k) model = build_model(conf.model, model_conf, dataset.num_users, dataset.num_items, device) logger.info(model_conf) logger.info(dataset) trainer = Trainer(dataset=dataset, model=model, evaluator=evaluator,
os.makedirs(exp_log_dir) exp_visual_dir = os.path.join(exp_dir, "visual") if not os.path.exists(exp_visual_dir): os.makedirs(exp_visual_dir) exp_ckpt_dir = os.path.join(exp_dir, "checkpoints") if not os.path.exists(exp_ckpt_dir): os.makedirs(exp_ckpt_dir) now_str = datetime.datetime.now().__str__().replace(' ', '_') writer_path = os.path.join(exp_visual_dir, now_str) writer = SummaryWriter(writer_path) logger_path = os.path.join(exp_log_dir, now_str + ".log") logger = Logger(logger_path).get_logger() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu logger.info("basic configuration settings: {}".format(basic_configs)) loss_function = nn.CrossEntropyLoss() network_configs=collections.OrderedDict() network_configs['coord_input_dim']= 2 network_configs['embed_dim']= 256 network_configs['feat_dict_size']= 103 network_configs['hidden_size']= 256
from json import loads import pytest from requests.api import get from requests.api import post from dao.MySQLHelper import MySQLHelper from utils.Config import Config from utils.Logger import Logger logger = Logger('TestEdu').getlog() file = './conf/config.ini' @pytest.fixture(scope='class') def get_db(): config = Config(file) my = MySQLHelper() my.host = config.get_value(file, 'mysql', 'host') my.user = config.get_value(file, 'mysql', 'user') my.password = config.get_value(file, 'mysql', 'password') my.database = config.get_value(file, 'mysql', 'database') my.port = int(config.get_value(file, 'mysql', 'port')) return my @pytest.fixture(scope='class') def init_config(): config = Config(file) configs = { 'url': config.get_value(file, 'base', 'url'),
# initialize replay buffer R_val = ReplayBuffer(MDP_STATE_SIZE, 1, BUFFER_SIZE) R = TrajectoryReplayBuffer(MDP_STATE_SIZE, 1, BUFFER_SIZE, STEPS) buf = R.LoadBuffer(OUT_DIR + BUFFER_FILE) buf_val = R_val.LoadBuffer(OUT_DIR + BUFFER_VAL) if buf: OBSERVATION_PHASE = 0 EXP_PROB = EPSILON populated = R.GetOccupency() print("Replay buffer loaded from disk, occupied: " + str(populated)) else: print("Creating new replay buffer") # initialize logger L = Logger() log_not_empty = L.Load(OUT_DIR + LOG_FILE) if log_not_empty: print("Log file loaded") else: ("Creating new log file") L.AddNewLog('network_left') L.AddNewLog('network_middle') L.AddNewLog('network_right') L.AddNewLog('policy_left') L.AddNewLog('policy_middle') L.AddNewLog('policy_right') L.AddNewLog('error') L.AddNewLog('total_reward') L.AddNewLog('estimated_value')
from utils.Logger import Logger from process.train_cus_model import Train if __name__ == '__main__': log = Logger('log/train_log_mymodel_res50.log', level='info') log.logger.info('Start our coursework!') train_arg = { 'original_dir': '/home/fanfu/newdisk/dataset/705/auair/', 'train_dir': '/home/fanfu/newdisk/dataset/705/auair/train/', 'test_dir': '/home/fanfu/newdisk/dataset/705/auair/test/', # 'train_dir': '/home/fanfu/newdisk/dataset/705/auair/test_auair_one (copy)/', # 'test_dir': '/home/fanfu/newdisk/dataset/705/auair/test_auair_one (copy)/', # 'train_dir': '/home/fanfu/newdisk/dataset/705/auair/auairdataset-master/examples/auair_subset/', # 'test_dir': '/home/fanfu/newdisk/dataset/705/auair/auairdataset-master/examples/auair_subset/', 'log': log, 'model': 'resnet50', # 'weight_save_path': 'model/first_model_weight_1.pth', 'weight_save_path': 'model/weight/fasterrcnn_resnet50_mymodel.pth', 'img_width': 1920, 'img_height': 1080, 'classes': [ "Background", "Human", "Car", "Truck", "Van", "Motorbike", "Bicycle", "Bus", "Trailer"
import matplotlib.pyplot as plt BASE = '/home/ayal/Documents/gym/Code/DQN_hockey/statistics/' files = ['log_DDQN200', 'log_DDQN1000', 'log_DDQN5000'] names = ['DDQN_200', 'DDQN_1000', 'DDQN_5000'] logs = [ 'network_left', 'network_middle', 'network_right', 'network_random', 'estimated_value' ] lognames = [ 'Puck left', 'Puck middle', 'Puck right', 'Puck random', 'Estimated value' ] readers = [] for i in xrange(len(files)): readers.append(Logger()) fname = BASE + files[i] readers[i].Load(BASE + files[i]) #left log = 'network_left' for log_index in xrange(len(logs)): log = logs[log_index] fig = plt.figure(log_index + 1) p1 = readers[0].GetLogByName(log) p1_up = readers[0].GetLogByName(log + '_up') p1_down = readers[0].GetLogByName(log + '_down') t1 = np.arange(1, len(p1) + 1) p2 = readers[1].GetLogByName(log)
# 'pre_train_weight_dir': 'model/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth', # # 'pre_train_weight_dir': 'model/first_model_weight_1.pth', # 'pre_train_self': False, # 'epoch_num': 10, # 'debug': False, # 'proportion': 0.8, # 'batch_size': 8, # 'learning_rate': 1e-3, # 'optimizer_mode': 'SGD' # } # # add Background to classes is for testing the human detection # train = Train(**train_arg) # train.train() ########################### baseline ########################################################### ########################### model 2 ########################################################### log = Logger('log/train_log_lrStep_model.log', level='info') log.logger.info('Start our coursework!') train_arg = { 'original_dir': '/home/fanfu/newdisk/dataset/705/auair/', 'train_dir': '/home/fanfu/newdisk/dataset/705/auair/train/', 'test_dir': '/home/fanfu/newdisk/dataset/705/auair/test/', # 'img_dir': '/home/fanfu/newdisk/dataset/705/auair/images', # 'label_dir': '/home/fanfu/newdisk/dataset/705/auair/annotations.json', # 'img_dir': '/home/fanfu/newdisk/dataset/705/auair/test_auair_one (copy)/images', # 'label_dir': '/home/fanfu/newdisk/dataset/705/auair/test_auair_one (copy)/annotations.json', # 'img_dir': '/home/fanfu/newdisk/dataset/705/auair/test_auair_one/images', # 'label_dir': '/home/fanfu/newdisk/dataset/705/auair/test_auair_one/annotations.json', # 'img_dir': '/home/fanfu/newdisk/dataset/705/auair/auairdataset-master/examples/auair_subset/images',
def train(): torch.manual_seed(args.seed) with open(args.arg, 'w') as f: for k,v in sorted(vars(args).items()): print("{0}: {1}".format(k,v), file = f) logger = Logger(model_id = args.unique_id, experiment_id = args.tb_experiment) print('Run ID:', args.unique_id) ''' Create model and load pretrained state ''' generator = Generator( channel = args.channels, filter_size = args.filter_size, scale = args.upscale_scale, training = True, improved_estimation = args.improved_estimation, detach_estimation = args.detach_estimation, interpolated_frames = args.interpolationsteps) if args.gan: in_channels = 0 in_channels += 784 * args.interpolationsteps * args.context in_channels += 8 * args.interpolationsteps * args.depth in_channels += 4 * args.interpolationsteps * args.flow in_channels += 3 * (args.interpolationsteps + 4) * args.spatial in_channels += 9 * args.interpolationsteps * args.temporal #TODO: 15 discriminator = Discriminator( in_channels = in_channels, training = True) # Save generator and discriminator when shutting down training killer.attach_generator(generator) if args.gan: killer.attach_discriminator(discriminator) if args.loss_perceptual: vgg = loss_perceptual.vgg19() vgg.eval() if args.use_cuda: vgg = vgg.cuda() if args.gradient_clipping > 0: tools.clip_gradients(model = generator, magnitude = args.gradient_clipping) if args.use_cuda: print("Turn the models into CUDA") generator = generator.cuda() if args.gan: discriminator = discriminator.cuda() ''' Load pretrained weights ''' if not args.pretrained_generator==None: tools.load_model_weights(model = generator, weights = args.pretrained_generator, use_cuda = args.use_cuda) if not args.pretrained_discriminator == None: tools.load_model_weights(model = discriminator, weights = args.pretrained_discriminator, use_cuda = args.use_cuda) if not args.pretrained_merger == None: tools.load_model_weights(model = generator.mergeNet, weights = args.pretrained_merger, use_cuda = args.use_cuda) if not args.pretrained_upscaler == None: tools.load_model_weights(model = generator.upscaleNet, weights = args.pretrained_upscaler, use_cuda = args.use_cuda) if args.perturb != 0: tools.perturb_weights(ctx = generator, p = args.perturb) tools.perturb_weights(ctx = discriminator, p = args.perturb) ''' Load Data ''' train_set, valid_set = datasets.UpscalingData( root = args.datasetPath, max_frame_size = (args.max_img_width, args.max_img_height), improved_estimation = args.improved_estimation, seqlength = args.interpolationsteps + 2, upscaling_scale = args.upscale_scale) train_loader = torch.utils.data.DataLoader( train_set, batch_size = args.batch_size, sampler = balancedsampler.RandomBalancedSampler(train_set, int(len(train_set) / args.batch_size )), num_workers = args.workers, pin_memory = True ) val_loader = torch.utils.data.DataLoader( valid_set, batch_size = args.batch_size, num_workers = args.workers, pin_memory = True ) print('{} samples found, {} train samples and {} test samples '.format(len(valid_set)+len(train_set), len(train_set), len(valid_set))) ''' Optimizers ''' print("train the interpolation net") # 1 step -> gradient from frame 1 # 2 steps-> gradient from frame 1 + gradient from frame 2 + backpropagation through recurrent frame 1 # ... => scale lr_scale = 2 / (args.interpolationsteps * (args.interpolationsteps + 1)) optimizer_g = torch.optim.Adamax([ {'params': generator.ctxNet.parameters(), 'lr': args.coe_ctx * args.lr_generator * lr_scale}, {'params': generator.depthNet.parameters(), 'lr': args.coe_depth * args.lr_generator * lr_scale}, {'params': generator.flownets.parameters(), 'lr': args.coe_flow * args.lr_generator * lr_scale}, {'params': generator.initScaleNets_filter.parameters(), 'lr': args.coe_filter * args.lr_generator * lr_scale}, {'params': generator.initScaleNets_filter1.parameters(), 'lr': args.coe_filter * args.lr_generator * lr_scale}, {'params': generator.initScaleNets_filter2.parameters(), 'lr': args.coe_filter * args.lr_generator * lr_scale}, {'params': generator.mergeNet.parameters(), 'lr': args.coe_merge * args.lr_generator * lr_scale}, {'params': generator.rectifyNet.parameters(), 'lr': args.lr_rectify * lr_scale}, {'params': generator.upscaleNet.parameters(), 'lr': args.lr_upscale * lr_scale} ], lr = args.lr_generator * lr_scale, betas = (0.9, 0.999), eps = 1e-8, weight_decay = args.weight_decay) if args.gan: optimizer_d = torch.optim.Adam( params = discriminator.parameters(), lr = args.lr_discriminator if args.loss_layer else args.lr_discriminator * 0.3, betas = (0.5, 0.999), eps = 1e-8, weight_decay = args.weight_decay ) if args.warmup_discriminator and args.warmup_boost: for param_group in optimizer_d.param_groups: param_group['lr'] = param_group['lr']*args.warmup_boost print('Discriminator learning rate boosted to', param_group['lr']) scheduler_g = lr_scheduler.ReduceLROnPlateau(optimizer_g, 'min', factor = args.factor, patience = args.patience, verbose = True) if args.gan: scheduler_d = lr_scheduler.ReduceLROnPlateau(optimizer_d, 'max', factor = 0.5, patience = 5, verbose = True) # Only used for kickstarting discriminator ''' Output loads of stuff ''' print("*********Start Training********") print("Generator LR is: "+ str(float(optimizer_g.param_groups[0]['lr']))) print("EPOCH is: "+ str(int(len(train_set) / args.batch_size ))) print("Num of EPOCH is: "+ str(args.numEpoch)) print("Num. of generator parameters:", tools.count_network_parameters(generator)) if hasattr(generator,'flownets'): print("Num. of flow model parameters:", tools.count_network_parameters(generator.flownets)) if hasattr(generator,'initScaleNets_occlusion'): print("Num. of initScaleNets_occlusion model parameters:", tools.count_network_parameters(generator.initScaleNets_occlusion) + tools.count_network_parameters(generator.initScaleNets_occlusion1) + tools.count_network_parameters(generator.initScaleNets_occlusion2)) if hasattr(generator,'initScaleNets_filter'): print("Num. of initScaleNets_filter model parameters:", tools.count_network_parameters(generator.initScaleNets_filter) + tools.count_network_parameters(generator.initScaleNets_filter1) + tools.count_network_parameters(generator.initScaleNets_filter2)) if hasattr(generator, 'ctxNet'): print("Num. of ctxNet model parameters:", tools.count_network_parameters(generator.ctxNet)) if hasattr(generator, 'depthNet'): print("Num. of depthNet model parameters:", tools.count_network_parameters(generator.depthNet)) if hasattr(generator,'rectifyNet'): print("Num. of rectifyNet model parameters:", tools.count_network_parameters(generator.rectifyNet)) if hasattr(generator, 'mergeNet'): print('Num. of merge network parameters:', tools.count_network_parameters(generator.mergeNet)) if hasattr(generator,'upscaleNet'): print("Num. of upscaleNet model parameters:", tools.count_network_parameters(generator.upscaleNet)) if args.gan: print("Num. of discriminator model parameters:", tools.count_network_parameters(discriminator)) ''' Define parameters and variables for stats ''' t_loss_charbonnier_sr = AverageMeter() t_loss_charbonnier_lr = AverageMeter() t_loss_perceptual = AverageMeter() t_loss_pingpong = AverageMeter() t_accuracy = RunAVGMeter(val = 0, size = args.warmup_lag) v_loss = AverageMeter() v_loss_charbonnier = AverageMeter() v_loss_perceptual = AverageMeter() v_loss_best = 10e10 dis_trained = 0 gen_trained = 0 # Epoch length training if args.lengthEpoch > len(train_set) or args.lengthEpoch < 1: epoch_length = int(len(train_set) / args.batch_size) else: epoch_length = int(args.lengthEpoch / args.batch_size) # Epoch length validation coe_val = (1-args.test_valid_split) if args.lengthEpoch * coe_val > len(valid_set) or args.lengthEpoch < 1: valid_length = int(len(valid_set) / args.batch_size) else: valid_length = int(args.lengthEpoch * coe_val / args.batch_size) # Reduce discriminator warmup time by not computing useless generator losses if args.warmup_discriminator: tmp_loss_perceptual = args.loss_perceptual tmp_loss_pingpong = args.loss_pingpong tmp_loss_layer = args.loss_layer args.loss_perceptual = False args.loss_pingpong = False args.loss_layer = False debug_root = os.path.join(os.getcwd(), args.debug_output_dir) def __step_train(i, hr, lr, epoch_length, epoch_dis, epoch_gen, dis_trained, gen_trained, sr_fadein, log_this_it, save_this_it): ''' 1 Input Handling: - high resolution frames to cuda - first and last frame of low resolution sequence to cuda - pad inputs with H or W < 128 ''' assert(len(hr) == len(lr) == (args.interpolationsteps + 2)) hr = [ Variable(frame.cuda() if args.use_cuda else frame, requires_grad = args.log_gradients) for frame in hr] lr = [ Variable(frame.cuda() if args.use_cuda else frame, requires_grad = args.log_gradients) for frame in lr] lr_start = lr[0] lr_end = lr[-1] ''' 2.1 Recurrent interpolation: - frames - charbonnier loss - perceptual loss - ping-pong loss ''' outputs_sr, outputs_lr = generator( frame_start = lr_start, frame_end = lr_end, low_memory = args.low_memory) loss_charbonnier_sr = loss_function.charbonnier_loss(output = outputs_sr, target = hr, epsilon = args.epsilon) loss_charbonnier_lr = loss_function.charbonnier_loss(output = outputs_lr[1:-1], target = lr[1:-1], epsilon = args.epsilon) if args.loss_perceptual: loss_perceptual = 0 for h, s in zip(hr, outputs_sr): vgg_real = vgg(x = h, normalize = True) vgg_fake = vgg(x = s, normalize = True) loss_perceptual += loss_function.cosine_similarity(vgg_real, vgg_fake) if args.loss_pingpong: #TODO: LR ping pong as well? outputs_sr_rev = generator(frame_start = lr_end, frame_end = lr_start, low_memory = args.low_memory)[0] outputs_sr_rev.reverse() loss_pingpong = loss_function.pingpong_loss(outputs_sr, outputs_sr_rev) ''' 2.2 Discriminator: - GAN loss - Layer loss ''' if args.gan: discriminator_input_real = generator.prepare_discriminator_inputs( sr = hr, frame0 = lr[0], frame1 = lr[-1], temporal = args.temporal, spatial = args.spatial, context = args.context, depth = args.depth, flow = args.flow) discriminator_input_fake = generator.prepare_discriminator_inputs( sr = outputs_sr, frame0 = lr[0], frame1 = lr[-1], temporal = args.temporal, spatial = args.spatial, context = args.context, depth = args.depth, flow = args.flow) score_real, layers_real = discriminator(discriminator_input_real.detach()) score_fake, _ = discriminator(discriminator_input_fake.detach()) score_gen , layers_gen = discriminator(discriminator_input_fake) loss_gen, loss_dis, loss_real, loss_fake = loss_function.gan_loss( score_real = score_real, score_fake = score_fake, score_gen = score_gen, epsilon = args.gan_epsilon, batchsize = args.batch_size) loss_layer = loss_function.layerloss(layers_real, layers_gen) if args.loss_layer else 0 ''' 3 Loss handling: - weighting - balance generator and discriminator power - backpropagation ''' # Combine losses loss_total = loss_charbonnier_lr * args.scale_lr / (args.scale_lr + args.scale_sr * sr_fadein) loss_total += loss_charbonnier_sr * args.scale_sr * sr_fadein / (args.scale_lr + args.scale_sr * sr_fadein) if args.loss_layer: loss_total += loss_layer * args.scale_layer * sr_fadein if args.loss_perceptual: loss_total += loss_perceptual * args.scale_perceptual * sr_fadein if args.loss_pingpong: loss_total += loss_pingpong * args.scale_pingpong * sr_fadein if args.gan: loss_total += loss_gen * args.scale_gan * sr_fadein # Scheduling if args.gan: gen_behindness = math.fabs(loss_gen.item()) - math.fabs(loss_real.item()) train_dis = False if args.freeze_dis else (True if args.freeze_gen else (gen_behindness < args.dis_threshold)) train_gen = False if args.freeze_gen else (True if args.freeze_dis else (gen_behindness > args.gen_threshold)) dis_trained += train_dis gen_trained += train_gen epoch_dis += train_dis epoch_gen += train_gen # Backpropagation optimizer_d.zero_grad() optimizer_g.zero_grad() if train_dis: loss_dis.backward(retain_graph = train_gen) # intermediate layers of discriminator used for layerloss for generator if train_gen: loss_total.backward() else: optimizer_g.zero_grad() loss_total.backward() train_gen = True ''' 4 Logging: - AverageMeters to print to screen - Losses to tensorboard - Gradients and weights to tensorboard - Save debugging frames to disk ''' if args.gan: t_acc_real, t_acc_fake = tools.accuracy(score_real = torch.mean(score_real).item(), score_fake = torch.mean(score_fake).item()) t_accuracy.update(val = 0.5*t_acc_real + 0.5*t_acc_fake, weight = args.batch_size) t_loss_charbonnier_sr.update(val = loss_charbonnier_sr.item(), n = args.batch_size) t_loss_charbonnier_lr.update(val = loss_charbonnier_lr.item(), n = args.batch_size) t_loss_perceptual.update(val = loss_perceptual.item() if args.loss_perceptual else 0, n = args.batch_size) t_loss_pingpong.update(val = loss_pingpong.item() if args.loss_pingpong else 0, n = args.batch_size) logger.log_scalars(tag = 'Generator', tag_value_dict = { 'total': loss_total.item(), 'charbonnier': loss_charbonnier_sr.item(), 'charbonnier_lr': loss_charbonnier_lr.item(), 'gan': loss_gen.item() if args.gan else 0, 'layer': loss_layer.item() if args.loss_layer else 0, 'pingpong': loss_pingpong.item() if args.loss_pingpong else 0, 'perceptual': loss_perceptual.item() if args.loss_perceptual else 0}, epoch = t, n_batch = i, num_batches = epoch_length) if args.gan: logger.log_scalars(tag = 'Discriminator/loss', tag_value_dict = { 'real': loss_real.item(), 'fake': loss_fake.item(), 'gen_behindness': gen_behindness}, epoch = t, n_batch = i, num_batches = epoch_length) logger.log_scalars(tag = 'Discriminator/scores', tag_value_dict = { 'real': torch.mean(score_real).item(), 'fake': torch.mean(score_fake).item()}, epoch = t, n_batch = i, num_batches = epoch_length) logger.log_scalars(tag = 'Discriminator/detection_performance', tag_value_dict = { 'real': t_acc_real, 'fake': t_acc_fake, 'avg': t_accuracy.avg}, epoch = t, n_batch = i, num_batches = epoch_length) logger.log_scalars(tag = 'Scheduling', tag_value_dict = { 'train_discriminator': 1 if train_dis else 0, 'train_generator': 1 if train_gen else 0, 'gen_behindness': gen_behindness}, epoch = t, n_batch = i, num_batches = epoch_length) logger.log_scalars(tag = 'Overview_trained', tag_value_dict = { 'generator': gen_trained, 'discriminator': dis_trained}, epoch = t, n_batch = i, num_batches = epoch_length) if args.log_gradients: if train_gen: logger.log_histogram(tag = 'weights/filter', values = tools.model_parameters(generator.initScaleNets_filter, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/filter', values = tools.model_parameters(generator.initScaleNets_filter, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'weights/filter1', values = tools.model_parameters(generator.initScaleNets_filter1, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/filter1', values = tools.model_parameters(generator.initScaleNets_filter1, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'weights/filter2', values = tools.model_parameters(generator.initScaleNets_filter2, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/filter2', values = tools.model_parameters(generator.initScaleNets_filter2, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'weights/ctxNet', values = tools.model_parameters(generator.ctxNet, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/ctxNet', values = tools.model_parameters(generator.ctxNet, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'weights/flownets', values = tools.model_parameters(generator.flownets, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/flownets', values = tools.model_parameters(generator.flownets, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'weights/depthNet', values = tools.model_parameters(generator.depthNet, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/depthNet', values = tools.model_parameters(generator.depthNet, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'weights/rectifyNet', values = tools.model_parameters(generator.rectifyNet, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/rectifyNet', values = tools.model_parameters(generator.rectifyNet, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'weights/mergeNet', values = tools.model_parameters(generator.mergeNet, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/mergeNet', values = tools.model_parameters(generator.mergeNet, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'weights/upscaleNet', values = tools.model_parameters(generator.upscaleNet, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/upscaleNet', values = tools.model_parameters(generator.upscaleNet, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) if args.gan and train_dis: logger.log_histogram(tag = 'weights/discriminator', values = tools.model_parameters(discriminator, 'weights'), epoch = t, n_batch = i, num_batches = epoch_length) logger.log_histogram(tag = 'gradients/discriminator', values = tools.model_parameters(discriminator, 'gradients'), epoch = t, n_batch = i, num_batches = epoch_length) # Print debugging frames if args.debug_output and (log_this_it or i == 0): if args.tb_debugframes: logger.save_images(tag = args.unique_id + '_' + str(t).zfill(3) + '_' + str(i).zfill(5) + '/real', image = torch.cat([ t for t in hr ], dim = 0), epoch = t, n_batch = i, num_batches = epoch_length) logger.save_images(tag = args.unique_id + '_' + str(t).zfill(3) + '_' + str(i).zfill(5) + '/fake', image = torch.cat([ t for t in outputs_sr ], dim = 0), epoch = t, n_batch = i, num_batches = epoch_length) tools.print_tensor( path = os.path.join(debug_root, args.unique_id + '_' + str(t).zfill(3) + '_' + str(i).zfill(5) + '_real.png'), img = tools.printable_tensor([ t.detach().cpu().numpy() for t in hr ])) tools.print_tensor( path = os.path.join(debug_root, args.unique_id + '_' + str(t).zfill(3) + '_' + str(i).zfill(5) + '_fake.png'), img = tools.printable_tensor([ t.detach().cpu().numpy() for t in outputs_sr ])) if args.loss_pingpong: if args.tb_debugframes: logger.save_images(tag = args.unique_id + '_' + str(t).zfill(3) + '_' + str(i).zfill(5) + '/rev', image = torch.cat([ t for t in outputs_sr_rev ], dim = 0), epoch = t, n_batch = i, num_batches = epoch_length) tools.print_tensor( path = os.path.join(debug_root, args.unique_id + '_' + str(t).zfill(3) + '_' + str(i).zfill(5) + '_rev.png'), img = tools.printable_tensor([ t.detach().cpu().numpy() for t in outputs_sr_rev ])) ''' 5 Finish: - optimizer step - save intermediate weights ''' if args.gan and train_dis: if args.gradient_scaling > 0: tools.rescale_gradients(model = discriminator, magnitude = args.gradient_scaling) optimizer_d.step() if train_gen: if args.gradient_scaling > 0: tools.rescale_gradients(model = generator, magnitude = args.gradient_scaling) optimizer_g.step() # Save intermediate weights if save_this_it: torch.save(generator.state_dict(), os.path.join(args.save_path, str(t).zfill(3) + '_' + str(i).zfill(6) + "_GEN.pth")) if args.gan: torch.save(discriminator.state_dict(), os.path.join(args.save_path, str(t).zfill(3) + '_' + str(i).zfill(6) + "_DIS.pth")) return epoch_dis, epoch_gen, dis_trained, gen_trained def __step_validate(i, hr, lr): hr = [ Variable(frame.cuda() if args.use_cuda else frame, requires_grad = False) for frame in hr] lr_start = Variable(lr[0].cuda() if args.use_cuda else lr[0], requires_grad = False) lr_end = Variable(lr[-1].cuda() if args.use_cuda else lr[-1], requires_grad = False) outputs_sr = generator(frame_start = lr_start, frame_end = lr_end) loss_charbonnier = loss_function.charbonnier_loss(output = outputs_sr, target = hr, epsilon = args.epsilon) if args.loss_perceptual: loss_perceptual = 0 for h, s in zip(hr, outputs_sr): vgg_real = vgg(x = h, normalize = True) vgg_fake = vgg(x = s, normalize = True) loss_perceptual += loss_function.cosine_similarity(vgg_real, vgg_fake) loss_total = loss_charbonnier * args.scale_sr if args.loss_perceptual: loss_total += loss_perceptual * args.scale_perceptual v_loss_charbonnier.update(val = loss_charbonnier.item(), n = args.batch_size) v_loss_perceptual.update(val = loss_perceptual.item() if args.loss_perceptual else 0, n = args.batch_size) v_loss.update(val = loss_total.item(), n = args.batch_size) warmup_samples = 0 for t in range(args.numEpoch + args.warmup_discriminator): # extra epoch for discriminator pretraining at the beginning ''' Training ''' generator = generator.train() if args.gan: discriminator = discriminator.train() epoch_dis = 0 epoch_gen = 0 if args.warmup_discriminator: args.freeze_gen = True for i, (hr, lr) in enumerate(train_loader): if i >= epoch_length: if not args.warmup_discriminator: break # Logging anything to screen this iteration? log_this_it = (i % max(1, int(epoch_length/args.debug_output_freq)) == 0 and i > 0) save_this_it = (i % max(1, args.save_interval) == 0 and args.save_interval > 0 and i > 0) epoch_dis, epoch_gen, dis_trained, gen_trained = __step_train( i = i, hr = hr, lr = lr, epoch_length = epoch_length, epoch_dis = epoch_dis, epoch_gen = epoch_gen, dis_trained = dis_trained, gen_trained = gen_trained, sr_fadein = min(1, t/max(1, args.sr_fadein)), log_this_it = log_this_it, save_this_it = save_this_it) if log_this_it: print('Epoch:', t, ' [', '{:6d}'.format(i),'\\', '{:6d}'.format(epoch_length), ']', \ 'Average charbonnier loss:', '{:7.5f}'.format(t_loss_charbonnier_sr.avg), \ '\tAverage perceptual loss:', '{:7.5f}'.format(t_loss_perceptual.avg), \ '\tdis:', '{:6d}'.format(epoch_dis), \ '\tgen:', '{:6d}'.format(epoch_gen), \ '\tdiscriminator accuracy:', '{:7.5f}'.format(t_accuracy.avg)) t_loss_charbonnier_sr.reset() t_loss_perceptual.reset() t_loss_pingpong.reset() epoch_dis = 0 epoch_gen = 0 if args.warmup_discriminator: warmup_samples += 1 if warmup_samples > args.warmup_lag: if warmup_samples % (epoch_length // 5) == 0: scheduler_d.step(t_accuracy.avg) if t_accuracy.avg > args.warmup_threshold: torch.save(discriminator.state_dict(), os.path.join(args.save_path, 'pretrain_dis_' + str(warmup_samples).zfill(5) + '.pth')) args.freeze_gen = False args.warmup_discriminator = False args.loss_perceptual = tmp_loss_perceptual args.loss_pingpong = tmp_loss_pingpong args.loss_layer = tmp_loss_layer logger.set_offset(warmup_samples) for param_group in optimizer_d.param_groups: param_group['lr'] = param_group['lr']/args.warmup_boost print('Warmup finished: discriminator pretrained for', warmup_samples, 'samples.') print('Start training generator.') break ''' Validation ''' print('-'*50) print('Start validation:', datetime.now().isoformat()) torch.save(generator.state_dict(), os.path.join(args.save_path, "GEN_" + str(t).zfill(3) + ".pth")) if args.gan: torch.save(discriminator.state_dict(), os.path.join(args.save_path, "DIS_" + str(t).zfill(3) + ".pth")) generator = generator.eval() with torch.no_grad(): for i, (hr, lr) in enumerate(val_loader): if i >= valid_length: break __step_validate(i = i, hr = hr, lr = lr) if v_loss.avg < v_loss_best: torch.save(generator.state_dict(), os.path.join(args.save_path, "best_GEN.pth")) logger.log_scalars(tag = 'Validation', tag_value_dict = { 'charbonnier': v_loss_charbonnier.avg, 'perceptual': v_loss_perceptual.avg, 'total': v_loss.avg}, epoch = 0, n_batch = t, num_batches = args.numEpoch) print('-'*50) print('Epoch:', t) print('Validation loss:', '{:7.5f}'.format(v_loss.avg)) print('Validation charbonnier loss:', '{:7.5f}'.format(v_loss_charbonnier.avg)) print('Validation perceptual loss:', '{:7.5f}'.format(v_loss_perceptual.avg)) if v_loss.avg < v_loss_best: print('Best weights updated.') print('-'*50) v_loss_best = v_loss.avg scheduler_g.step(v_loss.avg) v_loss.reset() v_loss_charbonnier.reset() v_loss_perceptual.reset() print('='*100) print('Training finished') print('='*100)
""" Paths """ out_model_dir = 'out_models' """ Settings """ parser = ArgumentParser() gpuid = parser.get_argument('gpuid') n_epoch = parser.get_argument('epoch') n_latent = parser.get_argument('dims') batchsize = parser.get_argument('batchsize') logger = Logger(n_epoch) serializer = Serializer(out_model_dir) computational_graph_generator = ComputationalGraph('cg') """ Data loading """ height = pickle.load(open('height.pkl')) height['data'] = height['data'].astype(np.float32) height['data'] /= np.max(height['data']) height['target'] = height['target'].astype(np.int32) N_train = 800 x_train, x_test = np.split(height['data'], [N_train]) y_train, y_test = np.split(height['target'], [N_train]) N_test = y_test.size
def train_net(config): os.environ["CUDA_VISIBLE_DEVICES"] = config['train']['gpu_id'] data_loader = DataLoader(config) train_loader = torch.utils.data.DataLoader( data_loader, batch_size=config['train']['batch_size'], shuffle=True, num_workers=config['train']['num_workers'], worker_init_fn = worker_init_fn, drop_last=True, pin_memory=False) start_epoch = 0 running_metric_binary = runningScore(2) if not (os.path.exists(config['train']['checkpoints'])): os.mkdir(config['train']['checkpoints']) checkpoints = os.path.join(config['train']['checkpoints'],"DB_%s_bs_%d_ep_%d" % (config['train']['backbone'], config['train']['batch_size'], config['train']['n_epoch'])) if not (os.path.exists(checkpoints)): os.mkdir(checkpoints) model = DBNet(config).cuda() criterion = L1BalanceCELoss() optimizer = torch.optim.SGD(model.parameters(), lr=config['train']['base_lr'], momentum=0.99, weight_decay=5e-4) if config['train']['restore']: print('Resuming from checkpoint.') assert os.path.isfile(config['train']['resume']), 'Error: no checkpoint directory found!' checkpoint = torch.load(config['train']['resume']) start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) log_write = Logger(os.path.join(checkpoints, 'log.txt'), title=config['train']['backbone'], resume=True) else: print('Training from scratch.') log_write = Logger(os.path.join(checkpoints,'log.txt'), title=config['train']['backbone']) log_write.set_names([' epoch', 'Total loss', ' Bce loss', 'Thresh loss', ' L1 loss', 'Binary Acc', 'Binary IoU', ' rescall',' precision',' hmean']) max_hmean = -1 for epoch in range(start_epoch,config['train']['n_epoch']): model.train() bce_loss_list = [] thresh_loss_list = [] l1_loss_list = [] total_loss_list = [] if(config['train']['decay_method']=='e_decay'): adjust_learning_rate_poly(config['train']['base_lr'], optimizer, epoch, max_epoch=config['train']['n_epoch'], factor=0.9) else: adjust_learning_rate(config, optimizer, epoch,config['train']['gama']) for batch_idx, (imgs, gts, gt_masks, thresh_maps, thresh_masks) in enumerate(train_loader): imgs = Variable(imgs.cuda()) gts = Variable(gts.cuda()) gt_masks = Variable(gt_masks.cuda()) thresh_maps = Variable(thresh_maps.cuda()) thresh_masks = Variable(thresh_masks.cuda()) batch = {} batch['gt'] = gts batch['mask'] = gt_masks batch['thresh_map'] = thresh_maps batch['thresh_mask'] = thresh_masks pre = model(imgs) loss, metrics = criterion(pre, batch) optimizer.zero_grad() loss.backward() optimizer.step() score_binary = cal_binary_score(pre['binary'], gts, gt_masks.unsqueeze(1), running_metric_binary) bce_loss_list.append(metrics['bce_loss'].item()) thresh_loss_list.append(metrics['thresh_loss'].item()) l1_loss_list.append(metrics['l1_loss'].item()) total_loss_list.append(loss.item()) if batch_idx % config['train']['show_step'] == 0: if(config['train']['print_format']=='linux'): headers = ['epoch/epochs','batch/batchs' ,'TotalLoss' ,'BceLoss',' ThreshLoss','L1Loss', 'Binary Acc','Binary IoU', 'Lr Rate'] show_item = [[str(epoch)+'/'+str(config['train']['n_epoch']), str(batch_idx + 1)+'/'+str(len(train_loader)), get_str(np.mean(total_loss_list)), get_str(np.mean(bce_loss_list)), get_str(np.mean(thresh_loss_list)), get_str(np.mean(l1_loss_list)), get_str(score_binary['Mean Acc']), get_str(score_binary['Mean IoU']), get_str(optimizer.param_groups[0]['lr']) ]] print_table(headers,show_item,type_str='train') else: output_log = '({epoch}/{epochs}/{batch}/{size}) | TotalLoss: {total_loss:.4f} | BceLoss: {bce_loss:.4f} | ThreshLoss: {thresh_loss: .4f} | L1Loss: {l1_loss: .4f} | Binary Acc: {bin_acc: .4f} | Binary IoU: {bin_iou: .4f} | Lr: {lr: .4f}'.format( epoch=epoch, epochs=config['train']['n_epoch'] , batch=batch_idx + 1, size=len(train_loader), total_loss=np.mean(total_loss_list), bce_loss=np.mean(bce_loss_list), thresh_loss=np.mean(thresh_loss_list), l1_loss=np.mean(l1_loss_list), bin_acc=score_binary['Mean Acc'], bin_iou=score_binary['Mean IoU'], lr=optimizer.param_groups[0]['lr'] ) print(output_log) sys.stdout.flush() if( epoch > config['train']['start_val_epoch']): result_dict = val(model,config) rescall,precision,hmean = result_dict['recall'],result_dict['precision'],result_dict['hmean'] print('epoch:',epoch,'rescall:',rescall,'precision:',precision,'hmean:',hmean) else: rescall = 0 precision = 0 hmean = 0 log_write.append([epoch, np.mean(total_loss_list), np.mean(bce_loss_list), np.mean(thresh_loss_list), np.mean(l1_loss_list), score_binary['Mean Acc'], score_binary['Mean IoU'], rescall,precision,hmean]) if(hmean > max_hmean and config['train']['start_val_epoch'] < config['train']['n_epoch']): max_hmean = hmean save_checkpoint({ 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'lr': config['train']['base_lr'], 'optimizer': optimizer.state_dict(), }, checkpoint=checkpoints,filename='best_model.pth.tar') save_checkpoint({ 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'lr': config['train']['base_lr'], 'optimizer': optimizer.state_dict(), }, checkpoint=checkpoints)
def main(config): # For fast training. cudnn.benchmark = True os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" if len(sys.argv) > 1: os.environ["CUDA_VISIBLE_DEVICES"] = config.gpu else: os.environ["CUDA_VISIBLE_DEVICES"] = "5" global device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') version = config.version beta1 = 0.5 beta2 = 0.999 loader = data_loader.get_loader( "/media/data2/laixc/AI_DATA/expression_transfer/face12/crop_face", "/media/data2/laixc/AI_DATA/expression_transfer/face12/points_face", config) G = ExpressionGenerater() D = RealFakeDiscriminator() #FEN = FeatureExtractNet() id_D = IdDiscriminator() kp_D = KeypointDiscriminator() points_G = LandMarksDetect() log_file = None ####### 载入预训练网络 ###### resume_iter = config.resume_iter ckpt_dir = "/media/data2/laixc/Facial_Expression_GAN/ckpt-{}".format( version) log = Logger(os.path.join(ckpt_dir, 'log.txt')) if os.path.exists(os.path.join(ckpt_dir, '{}-G.ckpt'.format(resume_iter))): G_path = os.path.join(ckpt_dir, '{}-G.ckpt'.format(resume_iter)) G.load_state_dict( torch.load(G_path, map_location=lambda storage, loc: storage)) D_path = os.path.join(ckpt_dir, '{}-D.ckpt'.format(resume_iter)) D.load_state_dict( torch.load(D_path, map_location=lambda storage, loc: storage)) IdD_path = os.path.join(ckpt_dir, '{}-idD.ckpt'.format(resume_iter)) id_D.load_state_dict( torch.load(IdD_path, map_location=lambda storage, loc: storage)) kp_D_path = os.path.join(ckpt_dir, '{}-kpD.ckpt'.format(resume_iter)) kp_D.load_state_dict( torch.load(kp_D_path, map_location=lambda storage, loc: storage)) points_G_path = os.path.join(ckpt_dir, '{}-pG.ckpt'.format(resume_iter)) points_G.load_state_dict( torch.load(points_G_path, map_location=lambda storage, loc: storage)) else: resume_iter = 0 ##### 训练face2keypoint #### points_G_optimizer = torch.optim.Adam(points_G.parameters(), lr=0.0001, betas=(0.5, 0.9)) kp_D_optimizer = torch.optim.Adam(kp_D.parameters(), lr=0.0001, betas=(0.5, 0.9)) G_optimizer = torch.optim.Adam(G.parameters(), lr=0.0001, betas=(0.5, 0.9)) D_optimizer = torch.optim.Adam(D.parameters(), lr=0.001, betas=(0.5, 0.9)) idD_optimizer = torch.optim.Adam(id_D.parameters(), lr=0.001, betas=(0.5, 0.9)) G.to(device) id_D.to(device) D.to(device) kp_D.to(device) points_G.to(device) #FEN.to(device) #FEN.eval() log.print(config) # Start training from scratch or resume training. start_iters = resume_iter trigger_rec = 1 data_iter = iter(loader) average_cnt = 0 average_g_loss = 0 # Start training. print('Start training...') for i in range(start_iters, 150000): # =================================================================================== # # 1. Preprocess input data # # =================================================================================== # #faces, origin_points = next(data_iter) #_, target_points = next(data_iter) try: faces, origin_points = next(data_iter) except StopIteration: data_iter = iter(loader) faces, origin_points = next(data_iter) rand_idx = torch.randperm(origin_points.size(0)) target_points = origin_points[rand_idx] target_faces = faces[rand_idx] faces = faces.to(device) target_faces = target_faces.to(device) origin_points = origin_points.to(device) target_points = target_points.to(device) # =================================================================================== # # 3. Train the discriminator # # =================================================================================== # # Real fake Dis real_loss = -torch.mean(D(faces)) # big for real faces_fake = G(faces, target_points) fake_loss = torch.mean(D(faces_fake)) # small for fake # Compute loss for gradient penalty. alpha = torch.rand(faces.size(0), 1, 1, 1).to(device) x_hat = (alpha * faces.data + (1 - alpha) * faces_fake.data).requires_grad_(True) out_src = D(x_hat) d_loss_gp = gradient_penalty(out_src, x_hat) lambda_gp = 10 Dis_loss = real_loss + fake_loss + lambda_gp * d_loss_gp D_optimizer.zero_grad() Dis_loss.backward() D_optimizer.step() # ID Dis id_real_loss = -torch.mean(id_D(faces, target_faces)) # big for real faces_fake = G(faces, target_points) id_fake_loss = torch.mean(id_D(faces, faces_fake)) # small for fake # Compute loss for gradient penalty. alpha = torch.rand(target_faces.size(0), 1, 1, 1).to(device) x_hat = (alpha * target_faces.data + (1 - alpha) * faces_fake.data).requires_grad_(True) out_src = id_D(faces, x_hat) id_d_loss_gp = gradient_penalty(out_src, x_hat) id_lambda_gp = 10 id_Dis_loss = id_real_loss + id_fake_loss + id_lambda_gp * id_d_loss_gp idD_optimizer.zero_grad() id_Dis_loss.backward() idD_optimizer.step() # Keypoints Dis kp_real_loss = -torch.mean(kp_D(target_faces, target_points)) # big for real faces_fake = G(faces, target_points) points_fake = points_G(faces_fake) kp_fake_loss = torch.mean(kp_D(target_faces, points_fake)) # small for fake # Compute loss for gradient penalty. alpha = torch.rand(target_faces.size(0), 1, 1, 1).to(device) x_hat = (alpha * target_points.data + (1 - alpha) * points_fake.data).requires_grad_(True) out_src = kp_D(target_faces, x_hat) kp_d_loss_gp = gradient_penalty(out_src, x_hat) kp_lambda_gp = 10 kp_Dis_loss = kp_real_loss + kp_fake_loss + kp_lambda_gp * kp_d_loss_gp kp_D_optimizer.zero_grad() kp_Dis_loss.backward() kp_D_optimizer.step() # if (i + 1) % 5 == 0: # print("iter {} - d_real_loss {:.2}, d_fake_loss {:.2}, d_loss_gp {:.2}".format(i,real_loss.item(), # fake_loss.item(), # lambda_gp * d_loss_gp # )) # =================================================================================== # # 3. Train the keypointsDetecter # # =================================================================================== # points_detect = points_G(faces) detecter_loss_clear = torch.mean( torch.abs(points_detect - origin_points)) detecter_loss = detecter_loss_clear points_G_optimizer.zero_grad() detecter_loss.backward() points_G_optimizer.step() # =================================================================================== # # 3. Train the generator # # =================================================================================== # n_critic = 4 if (i + 1) % n_critic == 0: while True: # Original-to-target domain. faces_fake = G(faces, target_points) predict_points = points_G(faces_fake) g_keypoints_loss = -torch.mean( kp_D(target_faces, predict_points)) g_fake_loss = -torch.mean(D(faces_fake)) # reconstructs = G(faces_fake, origin_points) # g_cycle_loss = torch.mean(torch.abs(reconstructs - faces)) g_id_loss = -torch.mean(id_D(faces, faces_fake)) l1_loss = torch.mean(torch.abs(faces_fake - target_faces)) # feature_loss = torch.mean(torch.abs(FEN(faces_fake) - FEN(target_faces))) # 轮流训练 # if (i+1) % 50 == 0: # trigger_rec = 1 - trigger_rec # print("trigger_rec : ", trigger_rec) lambda_rec = config.lambda_rec # 2 to 4 to 8 lambda_l1 = config.lambda_l1 lambda_keypoint = config.lambda_keypoint # 100 to 50 lambda_fake = config.lambda_fake lambda_id = config.lambda_id lambda_feature = config.lambda_feature g_loss = lambda_keypoint * g_keypoints_loss + lambda_fake * g_fake_loss \ + lambda_id * g_id_loss + lambda_l1 * l1_loss # + lambda_feature*feature_loss G_optimizer.zero_grad() g_loss.backward() G_optimizer.step() # 加权平均 if average_cnt == 0: average_g_loss = g_loss.item() else: average_g_loss = 4 / 5 * average_g_loss + 1 / 5 * g_loss.item( ) average_cnt = average_cnt + 1 if g_loss.item() < average_g_loss: print("<<< g_loss {} < average_g_loss {}".format( g_loss.item(), average_g_loss)) break else: print(">>> g_loss {} > average_g_loss {}".format( g_loss.item(), average_g_loss)) # Print out training information. if (i + 1) % 4 == 0: log.print( "iter {} - d_real_loss {:.2}, d_fake_loss {:.2}, d_loss_gp {:.2}, id_real_loss {:.2}, " "id_fake_loss {:.2}, id_loss_gp {:.2}, kp_real_loss {:.2}, " "kp_fake_loss {:.2}, kp_loss_gp {:.2} , g_keypoints_loss {:.2}, " "g_fake_loss {:.2}, g_id_loss {:.2}, L1_loss {:.2}".format( i, real_loss.item(), fake_loss.item(), lambda_gp * d_loss_gp, id_real_loss.item(), id_fake_loss.item(), id_lambda_gp * id_d_loss_gp, kp_real_loss.item(), kp_fake_loss.item(), kp_lambda_gp * kp_d_loss_gp, lambda_keypoint * g_keypoints_loss.item(), lambda_fake * g_fake_loss.item(), lambda_id * g_id_loss.item(), lambda_l1 * l1_loss)) sample_dir = "gan-sample-{}".format(version) if not os.path.isdir(sample_dir): os.mkdir(sample_dir) if (i + 1) % 24 == 0: with torch.no_grad(): target_point = target_points[0] fake_face = faces_fake[0] face = faces[0] #reconstruct = reconstructs[0] predict_point = predict_points[0] sample_path_face = os.path.join( sample_dir, '{}-image-face.jpg'.format(i + 1)) save_image(denorm(face.data.cpu()), sample_path_face) # sample_path_rec = os.path.join(sample_dir, '{}-image-reconstruct.jpg'.format(i + 1)) # save_image(denorm(reconstruct.data.cpu()), sample_path_rec) sample_path_fake = os.path.join( sample_dir, '{}-image-fake.jpg'.format(i + 1)) save_image(denorm(fake_face.data.cpu()), sample_path_fake) sample_path_target = os.path.join( sample_dir, '{}-image-target_point.jpg'.format(i + 1)) save_image(denorm(target_point.data.cpu()), sample_path_target) sample_path_predict_points = os.path.join( sample_dir, '{}-image-predict_point.jpg'.format(i + 1)) save_image(denorm(predict_point.data.cpu()), sample_path_predict_points) print('Saved real and fake images into {}...'.format( sample_path_face)) # Save model checkpoints. model_save_dir = "ckpt-{}".format(version) if (i + 1) % 1000 == 0: if not os.path.isdir(model_save_dir): os.mkdir(model_save_dir) point_G_path = os.path.join(model_save_dir, '{}-pG.ckpt'.format(i + 1)) torch.save(points_G.state_dict(), point_G_path) kp_D_path = os.path.join(model_save_dir, '{}-kpD.ckpt'.format(i + 1)) torch.save(kp_D.state_dict(), kp_D_path) G_path = os.path.join(model_save_dir, '{}-G.ckpt'.format(i + 1)) torch.save(G.state_dict(), G_path) D_path = os.path.join(model_save_dir, '{}-D.ckpt'.format(i + 1)) torch.save(D.state_dict(), D_path) idD_path = os.path.join(model_save_dir, '{}-idD.ckpt'.format(i + 1)) torch.save(id_D.state_dict(), idD_path) print('Saved model checkpoints into {}...'.format(model_save_dir))
import sys sys.path.append("..") from app import database from utils.ApiResponse import ApiResponse from utils.Logger import Logger from utils.hash import sha256, hash_id from model.User import User from model.Token import Token TOKEN_EXPIRATION_TIME = 60 * 60 * 6 logger = Logger() class TokenService(): @staticmethod def generateUserToken(user_id: str): """ Creates a token for a specific user. Removes any token previously created for the user. """ response = ApiResponse() user = User.query.filter_by(id=user_id).first() timestamp = time.time() timestamp_millis = int(round(timestamp * 1000)) token_ids = sha256(hash_id(timestamp_millis + randint(0, 9999))) token_value = sha256(hash_id(timestamp_millis) + str(uuid4()))
# -*- coding: utf-8 -*- from json import loads from requests.api import get, post from utils.Logger import Logger logger = Logger('TestEducation').getlog() class TestEducation(): def tes_login(self): url = 'http://192.168.1.5:8080/recruit.students/login/in' params = { 'account': 'admin', 'pwd': '660B8D2D5359FF6F94F8D3345698F88C' } cookies = {'JSESSIONID': '4C7F9900296CCF401EA453523988978F'} logger.info('url为:' + url) response = get(url, params, cookies=cookies) logger.info('我的状态码为%s' % response.status_code) # print(response) # print(response.text) # print(response.status_code) assert response.status_code == 200 def tes_addSchool(self): url = 'http://192.168.1.5:8080/recruit.students/school/manage/addSchoolInfo' data = { 'schoolName': '蔡包子4', 'listSchoolType[0].id': '2',
def main(config): # For fast training. cudnn.benchmark = True os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" if len(sys.argv) > 1: os.environ["CUDA_VISIBLE_DEVICES"] = config.gpu else: os.environ["CUDA_VISIBLE_DEVICES"] = "5" global device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') version = config.version beta1 = 0.5 beta2 = 0.999 loader = data_loader.get_loader( "/media/data2/laixc/AI_DATA/expression_transfer/face12/crop_face", "/media/data2/laixc/AI_DATA/expression_transfer/face12/points_face", config) points_G = LandMarksDetect() G = TwoPointBigConvExpressionGenerater() D = SNResRealFakeDiscriminator() FEN = FeatureExtractNet() id_D = SNResIdDiscriminator() ####### 载入预训练网络 ###### resume_iter = config.resume_iter ckpt_dir = "/media/data2/laixc/Facial_Expression_GAN/ckpt-{}".format( version) if not os.path.isdir(ckpt_dir): os.mkdir(ckpt_dir) log = Logger(os.path.join(ckpt_dir, 'log.txt')) if os.path.exists(os.path.join(ckpt_dir, '{}-G.ckpt'.format(resume_iter))): G_path = os.path.join(ckpt_dir, '{}-G.ckpt'.format(resume_iter)) G.load_state_dict( torch.load(G_path, map_location=lambda storage, loc: storage)) D_path = os.path.join(ckpt_dir, '{}-D.ckpt'.format(resume_iter)) D.load_state_dict( torch.load(D_path, map_location=lambda storage, loc: storage)) IdD_path = os.path.join(ckpt_dir, '{}-idD.ckpt'.format(resume_iter)) id_D.load_state_dict( torch.load(IdD_path, map_location=lambda storage, loc: storage)) points_G_path = os.path.join(ckpt_dir, '{}-pG.ckpt'.format(resume_iter)) points_G.load_state_dict( torch.load(points_G_path, map_location=lambda storage, loc: storage)) else: resume_iter = 0 ##### 训练face2keypoint #### points_G_optimizer = torch.optim.Adam(points_G.parameters(), lr=0.0001, betas=(0.5, 0.9)) G_optimizer = torch.optim.Adam(G.parameters(), lr=0.0001, betas=(0.5, 0.9)) D_optimizer = torch.optim.Adam(D.parameters(), lr=0.001, betas=(0.5, 0.9)) idD_optimizer = torch.optim.Adam(id_D.parameters(), lr=0.001, betas=(0.5, 0.9)) G.to(device) id_D.to(device) D.to(device) points_G.to(device) FEN.to(device) FEN.eval() log.print(config) # Start training from scratch or resume training. start_iters = resume_iter trigger_rec = 1 data_iter = iter(loader) # Start training. print('Start training...') for i in range(start_iters, 150000): # =================================================================================== # # 1. Preprocess input data # # =================================================================================== # #faces, origin_points = next(data_iter) #_, target_points = next(data_iter) try: rotate_faces, rotate_points, faces, origin_points = next(data_iter) except StopIteration: data_iter = iter(loader) rotate_faces, rotate_points, faces, origin_points = next(data_iter) rand_idx = torch.randperm(origin_points.size(0)) target_points = origin_points[rand_idx] target_faces = faces[rand_idx] faces = faces.to(device) rotate_faces = rotate_faces.to(device) rotate_points = rotate_points.to(device) target_faces = target_faces.to(device) origin_points = origin_points.to(device) target_points = target_points.to(device) # =================================================================================== # # 3. Train the discriminator # # =================================================================================== # # Real fake Dis real_loss = torch.mean(softplus(-D(faces))) # big for real faces_fake = G(rotate_faces, rotate_points, target_points) fake_loss = torch.mean(softplus(D(faces_fake))) # small for fake Dis_loss = real_loss + fake_loss D_optimizer.zero_grad() Dis_loss.backward() D_optimizer.step() # ID Dis id_real_loss = torch.mean( softplus(-id_D(faces, target_faces))) # big for real faces_fake = G(rotate_faces, rotate_points, target_points) id_fake_loss = torch.mean(softplus(id_D(faces, faces_fake))) # small for fake id_Dis_loss = id_real_loss + id_fake_loss idD_optimizer.zero_grad() id_Dis_loss.backward() idD_optimizer.step() # if (i + 1) % 5 == 0: # print("iter {} - d_real_loss {:.2}, d_fake_loss {:.2}, d_loss_gp {:.2}".format(i,real_loss.item(), # fake_loss.item(), # lambda_gp * d_loss_gp # )) # =================================================================================== # # 3. Train the keypointsDetecter # # =================================================================================== # points_detect = points_G(faces) detecter_loss_clear = torch.mean( torch.abs(points_detect - origin_points)) # faces_fake = G(faces, target_points) # reconstructs = G(faces_fake, origin_points) # detecter_loss_vague = torch.mean(torch.abs(points_G(reconstructs) - origin_points)) # # lambda_vague = 0.1 # detecter_loss = detecter_loss_clear + lambda_vague*detecter_loss_vague detecter_loss = detecter_loss_clear points_G_optimizer.zero_grad() detecter_loss.backward() points_G_optimizer.step() # =================================================================================== # # 3. Train the generator # # =================================================================================== # n_critic = 1 if (i + 1) % n_critic == 0: # Original-to-target domain. faces_fake = G(rotate_faces, rotate_points, target_points) predict_points = points_G(faces_fake) g_keypoints_loss = torch.mean( torch.abs(predict_points - target_points)) g_fake_loss = torch.mean(softplus(-D(faces_fake))) # reconstructs = G(faces_fake, origin_points) # g_cycle_loss = torch.mean(torch.abs(reconstructs - faces)) g_id_loss = torch.mean(softplus(-id_D(faces, faces_fake))) l1_loss = torch.mean(torch.abs(faces_fake - target_faces)) feature_loss = torch.mean( torch.abs(FEN(faces_fake) - FEN(target_faces))) # 轮流训练 # if (i+1) % 50 == 0: # trigger_rec = 1 - trigger_rec # print("trigger_rec : ", trigger_rec) lambda_rec = config.lambda_rec # 2 to 4 to 8 lambda_l1 = config.lambda_l1 lambda_keypoint = config.lambda_keypoint # 100 to 50 lambda_fake = config.lambda_fake lambda_id = config.lambda_id lambda_feature = config.lambda_feature g_loss = lambda_keypoint * g_keypoints_loss + lambda_fake*g_fake_loss \ + lambda_id * g_id_loss + lambda_l1 * l1_loss + lambda_feature*feature_loss G_optimizer.zero_grad() g_loss.backward() G_optimizer.step() # Print out training information. if (i + 1) % 4 == 0: log.print( "iter {} - d_real_loss {:.2}, d_fake_loss {:.2}, id_real_loss {:.2}, " "id_fake_loss {:.2} , g_keypoints_loss {:.2}, " "g_fake_loss {:.2}, g_id_loss {:.2}, L1_loss {:.2}, feature_loss {:.2}" .format(i, real_loss.item(), fake_loss.item(), id_real_loss.item(), id_fake_loss.item(), lambda_keypoint * g_keypoints_loss.item(), lambda_fake * g_fake_loss.item(), lambda_id * g_id_loss.item(), lambda_l1 * l1_loss, lambda_feature * feature_loss.item())) sample_dir = "gan-sample-{}".format(version) if not os.path.isdir(sample_dir): os.mkdir(sample_dir) if (i + 1) % 24 == 0: with torch.no_grad(): target_point = target_points[0] fake_face = faces_fake[0] #face = faces[0] rotate_face = rotate_faces[0] #reconstruct = reconstructs[0] predict_point = predict_points[0] rotate_point = rotate_points[0] sample_path_face = os.path.join( sample_dir, '{}-image-face.jpg'.format(i + 1)) save_image(denorm(rotate_face.data.cpu()), sample_path_face) # sample_path_rec = os.path.join(sample_dir, '{}-image-reconstruct.jpg'.format(i + 1)) # save_image(denorm(reconstruct.data.cpu()), sample_path_rec) sample_path_fake = os.path.join( sample_dir, '{}-image-fake.jpg'.format(i + 1)) save_image(denorm(fake_face.data.cpu()), sample_path_fake) sample_path_target = os.path.join( sample_dir, '{}-image-rotate_point.jpg'.format(i + 1)) save_image(denorm(rotate_point.data.cpu()), sample_path_target) sample_path_target = os.path.join( sample_dir, '{}-image-target_point.jpg'.format(i + 1)) save_image(denorm(target_point.data.cpu()), sample_path_target) sample_path_predict_points = os.path.join( sample_dir, '{}-image-predict_point.jpg'.format(i + 1)) save_image(denorm(predict_point.data.cpu()), sample_path_predict_points) print('Saved real and fake images into {}...'.format( sample_path_face)) # Save model checkpoints. model_save_dir = "ckpt-{}".format(version) if (i + 1) % 1000 == 0: if not os.path.isdir(model_save_dir): os.mkdir(model_save_dir) point_G_path = os.path.join(model_save_dir, '{}-pG.ckpt'.format(i + 1)) torch.save(points_G.state_dict(), point_G_path) G_path = os.path.join(model_save_dir, '{}-G.ckpt'.format(i + 1)) torch.save(G.state_dict(), G_path) D_path = os.path.join(model_save_dir, '{}-D.ckpt'.format(i + 1)) torch.save(D.state_dict(), D_path) idD_path = os.path.join(model_save_dir, '{}-idD.ckpt'.format(i + 1)) torch.save(id_D.state_dict(), idD_path) print('Saved model checkpoints into {}...'.format(model_save_dir))
#!/user/bin/python # coding=utf-8 from utils.Logger import Logger from torch.utils.tensorboard import SummaryWriter from constant import TMP_ROOT from os.path import join import sys import os import time # Create temporal directory for storing files. if os.path.isdir(TMP_ROOT) is not True: os.makedirs(TMP_ROOT) # Create temporal directory for saving checkpoints. if os.path.isdir(join(TMP_ROOT, 'checkpoints')) is not True: os.makedirs(join(TMP_ROOT, 'checkpoints')) logger = Logger( join( TMP_ROOT, 'log_%s.txt' % time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) ) ) sys.stdout = logger epoch_writer = SummaryWriter(log_dir=join(TMP_ROOT, 'tensorboard/epoch')) iteration_writer = SummaryWriter(log_dir=join(TMP_ROOT, 'tensorboard/iteration'))