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 finetune_loader = data_loader_augment.get_loader( "/media/data2/laixc/AI_DATA/expression_transfer/face12/crop_face", "/media/data2/laixc/AI_DATA/expression_transfer/face12/points_face", config) loader, first_loader = data_loader.get_loader( "/media/data2/laixc/AI_DATA/expression_transfer/face12/crop_sideface", "/media/data2/laixc/AI_DATA/expression_transfer/face12/points_sideface", config) points_G = LandMarksDetect() G = FastTrackExpressionGenerater() 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 finetune_iter = iter(finetune_loader) data_iter = iter(loader) first_data_iters = iter(first_loader) # Start training. print('Start training...') for i in range(start_iters, 150000): # =================================================================================== # # 1. Preprocess input data # # =================================================================================== # if i % 2 == 0: try: faces, origin_points = next(finetune_iter) except StopIteration: finetune_iter = iter(finetune_loader) faces, origin_points = next(finetune_iter) rand_idx = torch.randperm(origin_points.size(0)) target_points = origin_points[rand_idx] target_faces = faces[rand_idx] first_faces = faces.to(device) target_faces = target_faces.to(device) first_points = origin_points.to(device) target_points = target_points.to(device) else: try: first_faces, first_points = next(first_data_iters) except StopIteration: first_data_iters = iter(first_loader) first_faces, first_points = next(first_data_iters) try: target_faces, target_points = next(data_iter) except StopIteration: data_iter = iter(loader) target_faces, target_points = next(data_iter) target_faces = target_faces.to(device) first_faces = first_faces.to(device) target_points = target_points.to(device) first_points = first_points.to(device) # =================================================================================== # # 3. Train the discriminator # # =================================================================================== # # Real fake Dis real_loss = torch.mean(softplus(-D(target_faces))) # big for real faces_fake = G(first_faces, 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(first_faces, target_faces))) # big for real faces_fake = G(first_faces, target_points) id_fake_loss = torch.mean(softplus(id_D(first_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() # =================================================================================== # # 3. Train the keypointsDetecter # # =================================================================================== # points_detect = points_G(target_faces) detecter_loss_clear = torch.mean( torch.abs(points_detect - target_points)) 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(first_faces, 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(first_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))) 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 or i % 8 == 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 or i % 20 == 0: with torch.no_grad(): target_point = target_points[0] fake_face = faces_fake[0] face = first_faces[0] target_face = target_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_face = os.path.join( sample_dir, '{}-image-targetface.jpg'.format(i + 1)) save_image(denorm(target_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) 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))
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_augment.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 = ExpressionGenerater() D = RealFakeDiscriminator() FEN = FeatureExtractNet() ####### 载入预训练网络 ###### 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)) 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)) G.to(device) D.to(device) points_G.to(device) FEN.to(device) FEN.eval() # 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, faces, origin_points = next(data_iter) except StopIteration: data_iter = iter(loader) rotate_faces, 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) rotate_faces = rotate_faces.to(device) # =================================================================================== # # 3. Train the discriminator # # =================================================================================== # # Real fake Dis real_loss = - torch.mean(D(faces)) # big for real faces_fake = G(rotate_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() # 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 = 4 if (i + 1) % n_critic == 0: # Original-to-target domain. faces_fake = G(rotate_faces, target_points) predict_points = points_G(faces_fake) g_keypoints_loss = torch.mean(torch.abs(predict_points - target_points)) g_fake_loss = - torch.mean(D(faces_fake)) reconstructs = G(faces_fake, origin_points) g_cycle_loss = torch.mean(torch.abs(reconstructs - faces)) 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_feature = config.lambda_feature g_loss = lambda_keypoint * g_keypoints_loss + lambda_fake*g_fake_loss \ + lambda_rec * g_cycle_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}, d_loss_gp {:.2} , g_keypoints_loss {:.2}, " "g_fake_loss {:.2}, g_cycle_loss {:.2}, L1_loss {:.2}, detecter_loss {:.2}, feature_loss {:.2}".format(i, real_loss.item(), fake_loss.item(), lambda_gp * d_loss_gp , lambda_keypoint*g_keypoints_loss.item(), lambda_fake*g_fake_loss.item() , lambda_rec*g_cycle_loss.item(), lambda_l1 * l1_loss, detecter_loss.item(), 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] rotate_faces = rotate_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(rotate_faces.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_rec)) # 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) print('Saved model checkpoints into {}...'.format(model_save_dir))