def Learn_D(D_model, loss_criterion, loss_criterion_gan, optimizer_D, batch_image, generated, \ batch_id_label, batch_pose_label, batch_ones_label, batch_zeros_label, epoch, steps, Nd, args): real_output = D_model(batch_image) syn_output = D_model( generated.detach()) # .detach() をすることで Generatorまでの逆伝播計算省略 # id,真偽, pose それぞれのロスを計算 L_id = loss_criterion(real_output[:, :Nd], batch_id_label) L_gan = loss_criterion_gan(real_output[:, Nd], batch_ones_label) + loss_criterion_gan( syn_output[:, Nd], batch_zeros_label) L_pose = loss_criterion(real_output[:, Nd + 1:], batch_pose_label) d_loss = L_gan + L_id + L_pose d_loss.backward() optimizer_D.step() log_learning(epoch, steps, 'D', d_loss.data[0], args) # Discriminator の強さを判別 flag_D_strong = Is_D_strong(real_output, syn_output, batch_id_label, batch_pose_label, Nd) return flag_D_strong
def Learn_G(D_model, loss_criterion_gan, optimizer_G, batch_front_data, generated, batch_id_label, batch_ones_label, epoch, steps, args): syn_output = D_model(generated) L_sim = loss_criterion(syn_output, batch_front_data) L_gan = loss_criterion_gan(syn_output, batch_ones_label) g_loss = L_sim + L_gan g_loss.backward() optimizer_G.step() log_learning(epoch, steps, 'G', g_loss.data[0], args)
def Learn_D(D_model, loss_criterion_gan, optimizer_D, batch_front_data, generated, batch_id_label, batch_ones_label, batch_zeros_label, epoch, steps, args): real_output = D_model(batch_front_data) syn_output = D_model(generated.detach()) d_loss = loss_criterion_gan(real_output, batch_ones_label) + \ loss_criterion_gan(syn_output, batch_zeros_label) d_loss.backward() optimizer_D.step() log_learning(epoch, steps, 'D', d_loss.data[0], args)
def Learn_G(D_model, loss_criterion, loss_criterion_gan, optimizer_G ,generated, \ batch_id_label, batch_ones_label, pose_code_label, epoch, steps, Nd, args): syn_output = D_model(generated) # id についての出力と元画像のラベル, 真偽, poseについての出力と生成時に与えたposeコード の ロスを計算 L_id = loss_criterion(syn_output[:, :Nd], batch_id_label) L_gan = loss_criterion_gan(syn_output[:, Nd], batch_ones_label) L_pose = loss_criterion(syn_output[:, Nd + 1:], pose_code_label) g_loss = L_gan + L_id + L_pose g_loss.backward() optimizer_G.step() log_learning(epoch, steps, 'G', g_loss.data[0], args)
def Learn_G_re(D_model, loss_criterion, loss_criterion_gan, loss_cycle, loss_l2, optimizer_G ,generated, generated_cycle, \ batch_img, batch_id_label, batch_ones_label, pose_code_label, pose_code_cycle_label, epoch, steps, Nd, args): syn_output = D_model(generated) syn_output_cycle = D_model(generated_cycle) # id についての出力と元画像のラベル, 真偽, poseについての出力と生成時に与えたposeコード の ロスを計算 L_id = loss_criterion(syn_output[:, :Nd], batch_id_label) L_gan = loss_criterion_gan(syn_output[:, Nd], batch_ones_label) L_pose = loss_criterion(syn_output[:, Nd + 1:], pose_code_label) L_id_cycle = loss_criterion(syn_output_cycle[:, :Nd], batch_id_label) L_gan_cycle = loss_criterion_gan(syn_output_cycle[:, Nd], batch_ones_label) L_pose_cycle = loss_criterion(syn_output_cycle[:, Nd + 1:], pose_code_cycle_label) L_re = loss_cycle(generated_cycle, batch_img) # L_tv = (loss_l2(generated_cycle[])) g_loss = L_gan + L_id + L_pose + 0.1 * L_re + L_gan_cycle + L_id_cycle + L_pose_cycle g_loss.backward() optimizer_G.step() log_learning(epoch, steps, 'G_re', g_loss.data[0], args)
def Learn_G(D_model, loss_criterion, loss_criterion_gan, optimizer_G, generated, batch_id_label, batch_ones_label, pose_code_label, epoch, steps, Nd, args, generated_unique=None, batch_id_label_unique=None, pose_code_label_unique=None, minibatch_size_unique=None): syn_output = D_model(generated) # id についての出力と元画像のラベル, 真偽, poseについての出力と生成時に与えたposeコード の ロスを計算 L_id = loss_criterion(syn_output[:, :Nd], batch_id_label) L_gan = loss_criterion_gan(syn_output[:, Nd], batch_ones_label) L_pose = loss_criterion(syn_output[:, Nd + 1:], pose_code_label) g_loss = L_gan + L_id + L_pose if args.multi_DRGAN: syn_output_unique = D_model(generated_unique) L_id_unique = loss_criterion(syn_output_unique[:, :Nd], batch_id_label_unique) L_gan_unique = loss_criterion_gan( syn_output_unique[:, Nd], batch_ones_label[:minibatch_size_unique]) L_pose_unique = loss_criterion(syn_output_unique[:, Nd + 1:], pose_code_label_unique) g_loss = g_loss + L_gan_unique + L_id_unique + L_pose_unique g_loss.backward() optimizer_G.step() log_learning(epoch, steps, 'G', g_loss.data[0], args)
def Train(Model, args): Nd = args.Nd beta1_Adam = args.beta1 beta2_Adam = args.beta2 if args.cuda: Model.cuda() optimizer = optim.Adam(Model.parameters(), lr=args.lr, betas=(beta1_Adam, beta2_Adam)) #optimizer = optim.SGD(Model.parameters(), lr=args.lr) Model.train() steps = 0 CUDNN.benchmark = True for epoch in range(args.start_epoch, args.epochs + 1): if args.step_learning: adjust_learning_rate(optimizer, epoch, args) transformed_dataset = FaceIdPoseDataset(args.train_csv_file, transform=transforms.Compose([ transforms.Resize(256), transforms.RandomCrop(224), transforms.ToTensor() ])) dataloader = DataLoader(transformed_dataset, batch_size=args.Train_Batch, shuffle=True) for i, batch_data in enumerate(dataloader): Model.zero_grad() batch_image = torch.FloatTensor(batch_data[0].float()) batch_id_label = batch_data[2] if args.cuda: batch_image, batch_id_label = batch_image.cuda( ), batch_id_label.cuda() batch_image, batch_id_label = Variable(batch_image), Variable( batch_id_label) steps += 1 Prediction = Model(batch_image) Loss = Model.ID_Loss(Prediction, batch_id_label) Loss.backward() optimizer.step() log_learning(epoch, steps, 'VGG16_Model', args.lr, Loss.data, args) writer.add_scalar('Train/Train_Loss', Loss, steps) # Validation_Process(Model, epoch, writer, args) Validation_Process(Model, epoch, writer, args) if epoch % args.save_freq == 0: if not os.path.isdir(args.snapshot_dir): os.makedirs(args.snapshot_dir) save_path = os.path.join(args.snapshot_dir, 'epoch{}.pt'.format(epoch)) torch.save(Model.state_dict(), save_path) save_checkpoint( { 'epoch': epoch + 1, 'Model': Model.state_dict(), 'optimizer': optimizer.state_dict(), }, save_dir=os.path.join(args.snapshot_dir, 'epoch{}'.format(epoch))) # export scalar data to JSON for external processing writer.export_scalars_to_json("./all_scalars.json") writer.close()
def train_single_DRGAN(images, id_labels, pose_labels, Nd, Np, Nz, D_model, G_model, args, start_epoch=1): ''' input: images: image PATH list id_labels / pose_labels: id / pose numpy arr. (Not one hot) Nd/Np/Nz: # of id/pose/noise D_model: discriminator G_model: generator args: shell arg. start_epoch: starting epoch for training ''' if args.cuda: D_model.cuda() G_model.cuda() D_model.train() G_model.train() lr_Adam = args.lr beta1_Adam = args.beta1 beta2_Adam = args.beta2 rndcrop_size = args.rndcrop_train_img_size eps = 10**-300 # for safe logarithm REAL_LABEL = 0.9 image_size = len(images) epoch_time = np.ceil(image_size / args.batch_size).astype(int) optimizer_D = optim.Adam(D_model.parameters(), lr=lr_Adam, betas=(beta1_Adam, beta2_Adam)) optimizer_G = optim.Adam(G_model.parameters(), lr=lr_Adam, betas=(beta1_Adam, beta2_Adam)) loss_criterion = nn.CrossEntropyLoss() loss_log = [] steps = 0 flag_D_strong = False for epoch in range(start_epoch, args.epochs + 1): # Load augmented data (using img path) transformed_dataset = FaceIdPoseDataset2(images, id_labels, pose_labels, transform=transforms.Compose([ RandomCrop((rndcrop_size, rndcrop_size)) ]), img_size=args.train_img_size) dataloader = DataLoader(transformed_dataset, batch_size=args.batch_size, shuffle=True, num_workers=8) for i, batch_data in enumerate(dataloader): D_model.zero_grad() G_model.zero_grad() batch_image = torch.FloatTensor(batch_data[0].float()) batch_id_label = batch_data[1] batch_pose_label = batch_data[2] minibatch_size = len(batch_image) # ノイズと姿勢コードを生成 # fixed_noise = torch.FloatTensor(np.random.uniform(-1,1, (minibatch_size, Nz))) fixed_noise = torch.FloatTensor( np.random.standard_normal((minibatch_size, Nz))) tmp = torch.LongTensor(np.random.randint(Np, size=minibatch_size)) pose_code = one_hot(tmp, Np) # Condition 付に使用 pose_code_label = torch.LongTensor(tmp) # CrossEntropy 誤差に使用 if args.cuda: batch_image, batch_id_label, batch_pose_label = \ batch_image.cuda(), batch_id_label.cuda(), batch_pose_label.cuda() fixed_noise, pose_code, pose_code_label = \ fixed_noise.cuda(), pose_code.cuda(), pose_code_label.cuda() batch_image, batch_id_label, batch_pose_label = \ Variable(batch_image), Variable(batch_id_label), Variable(batch_pose_label) fixed_noise, pose_code, pose_code_label = \ Variable(fixed_noise), Variable(pose_code), Variable(pose_code_label) # Generatorでイメージ生成 generated = G_model(batch_image, pose_code, fixed_noise) steps += 1 # バッチ毎に交互に D と G の学習, Dが90%以上の精度の場合は 1:4の比率で学習 if flag_D_strong: if i % 5 == 0: # Discriminator の学習 real_output = D_model(batch_image) syn_output = D_model(generated.detach( )) # .detach() をすることで Generatorまでのパラメータを更新しない # id,真偽, pose それぞれのロスを計算 L_id = loss_criterion(real_output[:, :Nd], batch_id_label) if args.use_lsgan: L_gan = Variable.sum( (real_output[:, Nd] - REAL_LABEL)**2.0 + (syn_output[:, Nd] - 0.0)**2.0) / minibatch_size else: L_gan = Variable.sum( real_output[:, Nd].sigmoid().log() * -1 + (1 - syn_output[:, Nd].sigmoid()).log() * -1) / minibatch_size L_pose = loss_criterion(real_output[:, Nd + 1:], batch_pose_label) d_loss = L_gan + L_id + L_pose d_loss.backward() optimizer_D.step() log_learning(epoch, steps, 'D', d_loss.data[0], args) # Discriminator の強さを判別 flag_D_strong = Is_D_strong(real_output, syn_output, batch_id_label, batch_pose_label, Nd) else: # Generatorの学習 syn_output = D_model(generated) # id についての出力と元画像のラベル, 真偽, poseについての出力と生成時に与えたposeコード の ロスを計算 L_id = loss_criterion(syn_output[:, :Nd], batch_id_label) if args.use_lsgan: L_gan = Variable.sum((syn_output[:, Nd] - REAL_LABEL)** 2.0) / minibatch_size else: L_gan = Variable.sum( syn_output[:, Nd].sigmoid().clamp(min=eps).log() * -1) / minibatch_size L_pose = loss_criterion(syn_output[:, Nd + 1:], pose_code_label) g_loss = L_gan + L_id + L_pose g_loss.backward() optimizer_G.step() log_learning(epoch, steps, 'G', g_loss.data[0], args) else: if i % 2 == 0: # Discriminator の学習 real_output = D_model(batch_image) syn_output = D_model(generated.detach( )) # .detach() をすることでGeneratorのパラメータを更新しない # id,真偽, pose それぞれのロスを計算 L_id = loss_criterion(real_output[:, :Nd], batch_id_label) if args.use_lsgan: L_gan = Variable.sum( (real_output[:, Nd] - REAL_LABEL)**2.0 + (syn_output[:, Nd] - 0.0)**2.0) / minibatch_size else: L_gan = Variable.sum( real_output[:, Nd].sigmoid().log() * -1 + (1 - syn_output[:, Nd].sigmoid()).log() * -1) / minibatch_size L_pose = loss_criterion(real_output[:, Nd + 1:], batch_pose_label) d_loss = L_gan + L_id + L_pose d_loss.backward() optimizer_D.step() log_learning(epoch, steps, 'D', d_loss.data[0], args) # Discriminator の強さを判別 flag_D_strong = Is_D_strong(real_output, syn_output, batch_id_label, batch_pose_label, Nd) else: # Generatorの学習 syn_output = D_model(generated) # id についての出力と元画像のラベル, 真偽, poseについての出力と生成時に与えたposeコード の ロスを計算 L_id = loss_criterion(syn_output[:, :Nd], batch_id_label) if args.use_lsgan: L_gan = Variable.sum((syn_output[:, Nd] - REAL_LABEL)** 2.0) / minibatch_size else: L_gan = Variable.sum( syn_output[:, Nd].sigmoid().clamp(min=eps).log() * -1) / minibatch_size L_pose = loss_criterion(syn_output[:, Nd + 1:], pose_code_label) g_loss = L_gan + L_id + L_pose g_loss.backward() optimizer_G.step() log_learning(epoch, steps, 'G', g_loss.data[0], args) # エポック毎にロスの保存 loss_log.append([epoch, d_loss.data[0], g_loss.data[0]]) if epoch % args.save_freq == 0: # 各エポックで学習したモデルを保存 if not os.path.isdir(args.save_dir): os.makedirs(args.save_dir) save_path_D = os.path.join(args.save_dir, 'epoch{}_D.pt'.format(epoch)) torch.save(D_model, save_path_D) save_path_G = os.path.join(args.save_dir, 'epoch{}_G.pt'.format(epoch)) torch.save(G_model, save_path_G) # 最後のエポックの学習前に生成した画像を1枚保存(学習の確認用) save_generated_image = generated[0].cpu().data.numpy().transpose( 1, 2, 0) save_generated_image = np.squeeze(save_generated_image) # min~max -> 0~255 save_generated_image -= save_generated_image.min() save_generated_image = save_generated_image / save_generated_image.max( ) save_generated_image = save_generated_image * 255.0 # save_generated_image = (save_generated_image+1)/2.0 * 255. save_generated_image = save_generated_image[:, :, [ 2, 1, 0 ]] # convert from BGR to RGB save_path_image = os.path.join( args.save_dir, 'epoch{}_generatedimage.png'.format(epoch)) misc.imsave(save_path_image, save_generated_image.astype(np.uint8)) # 学習終了後に,全エポックでのロスの変化を画像として保存 loss_log = np.array(loss_log) plt.plot(loss_log[:, 1], label="Discriminative Loss") plt.plot(loss_log[:, 2], label="Generative Loss") plt.legend(loc='upper right') plt.xlabel("Epoch") plt.ylabel("Loss") filename = os.path.join(args.save_dir, 'Loss_log.png') plt.savefig(filename, bbox_inches='tight')
def train_single_fnm(D_model, G_model, C_model, args): writer = SummaryWriter() D_lr = args.lr G_lr = args.lr beta1_Adam = args.beta1 beta2_Adam = args.beta2 if args.cuda: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if args.cuda: D_model.to(device) G_model.to(device) C_model.to(device) optimizer_D = optim.Adam(D_model.parameters(), lr=D_lr, betas=(beta1_Adam, beta2_Adam), weight_decay=args.lambda_reg) optimizer_G = optim.Adam(G_model.parameters(), lr=G_lr, betas=(beta1_Adam, beta2_Adam), weight_decay=args.lambda_reg) if args.resume: checkpoint = torch.load(args.resume) optimizer_D.load_state_dict(checkpoint['optimizer_D']) optimizer_G.load_state_dict(checkpoint['optimizer_G']) steps = 0 CUDNN.benchmark = True for epoch in range(args.start_epoch, args.epochs + 1): D_model.train() G_model.train() C_model.eval() # Load augmented data profile_dataset = FaceIdPoseDataset( args.profile_list, args.data_place, transform=transforms.Compose([ torchvision.transforms.Resize(250), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])) front_dataset = FaceIdPoseDataset( args.front_list, args.data_place, transform=transforms.Compose([ torchvision.transforms.Resize(224), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])) profile_dataloader = DataLoader(profile_dataset, batch_size=args.batch_size, shuffle=True) #, num_workers=6) front_dataloader = DataLoader(front_dataset, batch_size=args.batch_size, shuffle=True) # , num_workers=6) for idx, _ in enumerate(profile_dataloader): batch_profile, imageName_profile = next(iter(profile_dataloader)) batch_front, imageName_front = next(iter(front_dataloader)) batch_profile = ((batch_profile + 1) * 127.5).to(device) batch_front = ((batch_front + 1) * 127.5).to(device) steps += 1 enable_gradients(D_model) disable_gradients(C_model) disable_gradients(G_model) if steps < 25 and epoch == 1: critic = 25 else: critic = args.num_critic_D for _ in range(0, critic): D_model.zero_grad() # Create Encoder Feature Map / Get the Real images' Features _, Front_FeaMap = C_model(batch_front) _, Profile_FeaMap = C_model(batch_profile) gen_f = G_model(Front_FeaMap) gen_p = G_model(Profile_FeaMap) # Mapping to single unit by using Discriminator syn_f_gan = D_model(gen_f) syn_p_gan = D_model(gen_p) real_gan = D_model(batch_front) # Gradient Penalty gp_alpha = torch.FloatTensor(batch_front.size()[0], 1, 1, 1).to(device) gp_alpha.uniform_() interpolates = gen_p.data * gp_alpha + ( 1 - gp_alpha) * batch_front.data interpolates = interpolates.to( device).requires_grad_() # requires_grad_() 開啟張量 Loss, Wdis, GP = D_model.CriticWithGP_Loss( syn_f_gan, syn_p_gan, real_gan, interpolates) L_D = Loss L_D.backward() optimizer_D.step() writer.add_scalar('Discriminator/Gradient-Penalty', GP, steps) writer.add_scalar('Discriminator/Wasserstein-Distance', Wdis, steps) writer.add_scalar('Discriminator/D-LOSS', Loss, steps) log_learning(epoch, steps, 'D', D_lr, L_D.data, args) enable_gradients(G_model) disable_gradients(D_model) disable_gradients(C_model) for _ in range(0, args.num_critic_G): G_model.zero_grad() """Loss Functions 1. Pixel-Wise Loss: front-to-front reconstruct 2. Perceptual Loss: Feature distance on space of pretrined face model 3. Regulation Loss: L2 weight regulation (Aleady included in nn.Adam) 4. Adversarial Loss: Wasserstein Distance 5. Symmetric Loss: NOT APPLY 6. Drift Loss: NOT APPLY 7. Grade Penalty Loss: Grade penalty for Discriminator """ # Create Encoder Feature Map / Get the Real images' Features Front_Fea, Front_FeaMap = C_model(batch_front) Profile_Fea, Profile_FeaMap = C_model(batch_profile) # Synthesized image / Get the Fake images' Features gen_f = G_model(Front_FeaMap) gen_p = G_model(Profile_FeaMap) Front_Syn_Fea, _ = C_model(gen_f) Profile_Syn_Fea, _ = C_model(gen_p) # Mapping to single unit by using Discriminator syn_f_gan = D_model(gen_f) syn_p_gan = D_model(gen_p) # Frontalization Loss: L1-Norm L1 = G_model.L1Loss(gen_f, batch_front) #(input, target) # Feature Loss: Cosine-Norm / L2-Norm L2 = G_model.L2Loss(Front_Syn_Fea, Front_Fea, Profile_Syn_Fea, Profile_Fea) # Adversarial Loss L_Gen = G_model.GLoss(syn_f_gan, syn_p_gan) # L2 Regulation Loss (L2 regularization on the parameters of the model is already included in most optimizers) L_G = args.lambda_l1 * L1 + args.lambda_fea * L2 + args.lambda_gan * L_Gen L_G.backward() optimizer_G.step() writer.add_scalar('Generator/Pixel-Wise-Loss', L1, steps) writer.add_scalar('Generator/Perceptual-Loss', L2, steps) writer.add_scalar('Generator/Adversarial Loss', L_Gen, steps) writer.add_scalar('Generator/G-LOSS', L_Gen, steps) log_learning(epoch, steps, 'G', G_lr, L_G.data, args) if steps % 500 == 0: x_r = vutils.make_grid(batch_front, normalize=True, scale_each=True) y_r = vutils.make_grid(batch_profile, normalize=True, scale_each=True) x_f = vutils.make_grid(gen_f, normalize=True, scale_each=True) y_f = vutils.make_grid(gen_p, normalize=True, scale_each=True) writer.add_image('Image/Front-Real', x_r, steps) writer.add_image('Image/Front-Generated', x_f, steps) writer.add_image('Image/Profile-Real', y_r, steps) writer.add_image('Image/Profile-Generated', y_f, steps) save_path_image = os.path.join( args.snapshot_dir, 'epoch{}_FrontInput.jpg'.format(epoch)) torchvision.utils.save_image(batch_front, save_path_image, normalize=True, scale_each=True) save_path_image = os.path.join( args.snapshot_dir, 'epoch{}_FrontSynthesized.jpg'.format(epoch)) torchvision.utils.save_image(gen_f, save_path_image, normalize=True, scale_each=True) save_path_image = os.path.join( args.snapshot_dir, 'epoch{}_ProfileInput.jpg'.format(epoch)) torchvision.utils.save_image(batch_profile, save_path_image, normalize=True, scale_each=True) save_path_image = os.path.join( args.snapshot_dir, 'epoch{}_ProfileSynthesized.jpg'.format(epoch)) torchvision.utils.save_image(gen_p, save_path_image, normalize=True, scale_each=True) if epoch % args.save_freq == 0: if not os.path.isdir(args.snapshot_dir): os.makedirs(args.snapshot_dir) save_path_D = os.path.join(args.snapshot_dir, 'epoch{}_D.pt'.format(epoch)) torch.save(D_model.state_dict(), save_path_D) save_path_G = os.path.join(args.snapshot_dir, 'epoch{}_G.pt'.format(epoch)) torch.save(G_model.state_dict(), save_path_G) save_checkpoint( { 'epoch': epoch + 1, 'D_model': D_model.state_dict(), 'optimizer_D': optimizer_D.state_dict(), 'G_model': G_model.state_dict(), 'optimizer_G': optimizer_G.state_dict(), }, save_dir=os.path.join(args.snapshot_dir, 'epoch{}'.format(epoch))) # export scalar data to JSON for external processing writer.export_scalars_to_json("./all_scalars.json") writer.close()
def train_multiple_DRGAN(images_whole, id_labels_whole, pose_labels_whole, Nd, Np, Nz, D_model, G_model, args): if args.cuda: D_model.cuda() G_model.cuda() D_model.train() G_model.train() lr_Adam = args.lr beta1_Adam = args.beta1 beta2_Adam = args.beta2 eps = 10**-300 optimizer_D = optim.Adam(D_model.parameters(), lr=lr_Adam, betas=(beta1_Adam, beta2_Adam)) optimizer_G = optim.Adam(G_model.parameters(), lr=lr_Adam, betas=(beta1_Adam, beta2_Adam)) loss_criterion = nn.CrossEntropyLoss() loss_log = [] steps = 0 flag_D_strong = False for epoch in range(1, args.epochs + 1): images, id_labels, pose_labels = create_multiDR_GAN_traindata(images_whole,\ id_labels_whole, pose_labels_whole, args) image_size = images.shape[0] epoch_time = np.ceil(image_size / args.batch_size).astype(int) pdb.set_trace() for i in range(epoch_time): D_model.zero_grad() G_model.zero_grad() start = i * args.batch_size end = start + args.batch_size batch_image = torch.FloatTensor(images[start:end]) batch_id_label = torch.LongTensor(id_labels[start:end]) batch_id_label_unique = torch.LongTensor( batch_id_label[::args.images_perID]) batch_pose_label = torch.LongTensor(pose_labels[start:end]) minibatch_size = len(batch_image) minibatch_size_unique = len(batch_image) // args.images_perID # 特徴量をまとめた場合とそれぞれ用いた場合の ノイズと姿勢コードを生成 # それぞれ用いた場合 syn_id_label = torch.LongTensor( Nd * np.ones(minibatch_size).astype(int)) fixed_noise = torch.FloatTensor( np.random.uniform(-1, 1, (minibatch_size, Nz))) tmp = torch.LongTensor(np.random.randint(Np, size=minibatch_size)) pose_code = one_hot(tmp, Np) # Condition 付に使用 pose_code_label = torch.LongTensor(tmp) # CrossEntropy 誤差に使用 # 同一人物の特徴量をまとめた場合 syn_id_label_unique = torch.LongTensor( Nd * np.ones(minibatch_size_unique).astype(int)) fixed_noise_unique = torch.FloatTensor( np.random.uniform(-1, 1, (minibatch_size_unique, Nz))) tmp = torch.LongTensor( np.random.randint(Np, size=minibatch_size_unique)) pose_code_unique = one_hot(tmp, Np) # Condition 付に使用 pose_code_label_unique = torch.LongTensor( tmp) # CrossEntropy 誤差に使用 if args.cuda: batch_image, batch_id_label, batch_pose_label, syn_id_label, \ fixed_noise, pose_code, pose_code_label = \ batch_image.cuda(), batch_id_label.cuda(), batch_pose_label.cuda(), syn_id_label.cuda(), \ fixed_noise.cuda(), pose_code.cuda(), pose_code_label.cuda() batch_id_label_unique, syn_id_label_unique, \ fixed_noise_unique, pose_code_unique, pose_code_label_unique = \ batch_id_label_unique.cuda(), syn_id_label_unique.cuda(), \ fixed_noise_unique.cuda(), pose_code_unique.cuda(), pose_code_label_unique.cuda() batch_image, batch_id_label, batch_pose_label, syn_id_label, \ fixed_noise, pose_code, pose_code_label = \ Variable(batch_image), Variable(batch_id_label), Variable(batch_pose_label), Variable(syn_id_label), \ Variable(fixed_noise), Variable(pose_code), Variable(pose_code_label) batch_id_label_unique, syn_id_label_unique, \ fixed_noise_unique, pose_code_unique, pose_code_label_unique = \ Variable(batch_id_label_unique), Variable(syn_id_label_unique), \ Variable(fixed_noise_unique), Variable(pose_code_unique), Variable(pose_code_label_unique) # Generatorでイメージ生成 # 個々の画像特徴量からそれぞれ画像を生成した場合 generated = G_model(batch_image, pose_code, fixed_noise, single=True) # 同一人物の画像特徴量を一つにまとめた場合 generated_unique = G_model(batch_image, pose_code_unique, fixed_noise_unique) steps += 1 # バッチ毎に交互に D と G の学習, Dが90%以上の精度の場合は 1:4の比率で学習 if flag_D_strong: if i % 5 == 0: # Discriminator の学習 real_output = D_model(batch_image) syn_output = D_model(generated.detach( )) # .detach() をすることでGeneratorのパラメータを更新しない # id,真偽, pose それぞれのロスを計算 L_id = loss_criterion(real_output[:, :Nd], batch_id_label) L_gan = Variable.sum( real_output[:, Nd].sigmoid().calmp(min=eps).log() * -1 + (1 - syn_output[:, Nd].sigmoid()).clamp( min=eps).log() * -1) / minibatch_size L_pose = loss_criterion(real_output[:, Nd + 1:], batch_pose_label) d_loss = L_gan + L_id + L_pose d_loss.backward() optimizer_D.step() log_learning(epoch, steps, 'D', d_loss.data[0], args) # Discriminator の強さを判別 flag_D_strong = Is_D_strong(real_output, syn_output, batch_id_label, batch_pose_label, Nd) else: # Generatorの学習 syn_output = D_model(generated) syn_output_unique = D_model(generated_unique) # id についての出力と元画像のラベル, 真偽, poseについての出力と生成時に与えたposeコード の ロスを計算 L_id = loss_criterion(syn_output[:, :Nd], batch_id_label) L_gan = Variable.sum( syn_output[:, Nd].sigmoid().clamp(min=eps).log() * -1) / minibatch_size L_pose = loss_criterion(syn_output[:, Nd + 1:], pose_code_label) L_id_unique = loss_criterion(syn_output_unique[:, :Nd], batch_id_label_unique) L_gan_unique = Variable.sum( syn_output_unique[:, Nd].sigmoid().clamp( min=eps).log() * -1) / minibatch_size_unique L_pose_unique = loss_criterion( syn_output_unique[:, Nd + 1:], pose_code_label_unique) g_loss = L_gan + L_id + L_pose + L_gan_unique + L_id_unique + L_pose_unique g_loss.backward() optimizer_G.step() log_learning(epoch, steps, 'G', g_loss.data[0], args) else: if i % 2 == 0: # Discriminator の学習 real_output = D_model(batch_image) syn_output = D_model(generated.detach( )) # .detach() をすることでGeneratorのパラメータを更新しない # id,真偽, pose それぞれのロスを計算 L_id = loss_criterion(real_output[:, :Nd], batch_id_label) L_gan = Variable.sum( real_output[:, Nd].sigmoid().clamp(min=eps).log() * -1 + (1 - syn_output[:, Nd].sigmoid()).clamp( min=eps).log() * -1) / minibatch_size L_pose = loss_criterion(real_output[:, Nd + 1:], batch_pose_label) d_loss = L_gan + L_id + L_pose d_loss.backward() optimizer_D.step() log_learning(epoch, steps, 'D', d_loss.data[0], args) # Discriminator の強さを判別 flag_D_strong = Is_D_strong(real_output, syn_output, batch_id_label, batch_pose_label, Nd) else: # Generatorの学習 syn_output = D_model(generated) syn_output_unique = D_model(generated_unique) # id についての出力と元画像のラベル, 真偽, poseについての出力と生成時に与えたposeコード の ロスを計算 L_id = loss_criterion(syn_output[:, :Nd], batch_id_label) L_gan = Variable.sum( syn_output[:, Nd].sigmoid().clamp(min=eps).log() * -1) / minibatch_size L_pose = loss_criterion(syn_output[:, Nd + 1:], pose_code_label) L_id_unique = loss_criterion(syn_output_unique[:, :Nd], batch_id_label_unique) L_gan_unique = Variable.sum( syn_output_unique[:, Nd].sigmoid().clamp( min=eps).log() * -1) / minibatch_size_unique L_pose_unique = loss_criterion( syn_output_unique[:, Nd + 1:], pose_code_label_unique) g_loss = L_gan + L_id + L_pose + L_gan_unique + L_id_unique + L_pose_unique g_loss.backward() optimizer_G.step() log_learning(epoch, steps, 'G', g_loss.data[0], args) # エポック毎にロスの保存 loss_log.append([epoch, d_loss.data[0], g_loss.data[0]]) if epoch % args.save_freq == 0: # 各エポックで学習したモデルを保存 if not os.path.isdir(args.save_dir): os.makedirs(args.save_dir) save_path_D = os.path.join(args.save_dir, 'epoch{}_D.pt'.format(epoch)) torch.save(D_model, save_path_D) save_path_G = os.path.join(args.save_dir, 'epoch{}_G.pt'.format(epoch)) torch.save(G_model, save_path_G) # 最後のエポックの学習前に生成した画像を1枚保存(学習の確認用) save_generated_image = generated[0].cpu().data.numpy().transpose( 1, 2, 0) save_generated_image = np.squeeze(save_generated_image) save_generated_image = (save_generated_image + 1) / 2.0 * 255. save_generated_image = save_generated_image[:, :, [ 2, 1, 0 ]] # convert from BGR to RGB save_path_image = os.path.join( args.save_dir, 'epoch{}_generatedimage.jpg'.format(epoch)) misc.imsave(save_path_image, save_generated_image.astype(np.uint8)) # 学習終了後に,全エポックでのロスの変化を画像として保存 loss_log = np.array(loss_log) plt.plot(loss_log[:, 1], label="Discriminative Loss") plt.plot(loss_log[:, 2], label="Generative Loss") plt.legend(loc='upper right') plt.xlabel("Epoch") plt.ylabel("Loss") filename = os.path.join(args.save_dir, 'Loss_log.png') plt.savefig(filename, bbox_inches='tight')
def Train(Model, args): writer = SummaryWriter() beta1_Adam = args.beta1 beta2_Adam = args.beta2 if args.cuda: Model.cuda() #optimizer = optim.Adam(Model.parameters(), lr=args.lr, betas=(beta1_Adam, beta2_Adam)) optimizer = optim.SGD(Model.parameters(), lr=args.lr) if args.resume: checkpoint = torch.load(args.resume) optimizer.load_state_dict(checkpoint['optimizer']) Model.train() steps = 0 #loss_criterion_Angular = AngleLoss().cuda() CUDNN.benchmark = True if args.stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) if args.dynamic_lr == True: scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=3000, verbose=False, threshold=0.00001, threshold_mode='rel', cooldown=2000, min_lr=0, eps=1e-08) for epoch in range(args.start_epoch, args.epochs+1): #if epoch==3: #optimizer = optim.SGD(Model.parameters(), lr=args.lr) # Every args.lr_step, changes learning rate by multipling args.lr_decay #adjust_learning_rate(optimizer, epoch, args) # Load augmented data #transformed_dataset = FaceIdPoseDataset(args.train_csv_file, args.data_place, #transform = transforms.Compose([Resize((256,256)), RandomCrop((224,224))])) #for ResNet256x256->224x224 for VGG110x110->96x96 # transformed_dataset = FaceIdPoseDataset(args.train_csv_file, args.data_place, # transforms.Compose([transforms.Resize(256), transforms.RandomCrop(224),transforms.ToTensor()])) # for ResNet256x256->224x224 for VGG110x110->96x96 transformed_dataset = FaceIdPoseDataset(args.train_csv_file, args.data_place,transforms.Compose([transforms.Resize(256), transforms.RandomCrop(224), transforms.ToTensor() ])) # for ResNet256x256->224x224 for VGG110x110->96x96 dataloader = DataLoader(transformed_dataset, batch_size=args.Train_Batch, shuffle=True, num_workers=8) if args.stepsize > 0: scheduler.step() for i, batch_data in enumerate(dataloader): # backward() function accumulates gradients, however we don't want to mix up gradients between minibatches optimizer.zero_grad() batch_image = torch.FloatTensor(batch_data[0].float()) batch_id_label = batch_data[2] if args.cuda: batch_image, batch_id_label = batch_image.cuda(), batch_id_label.cuda() batch_image, batch_id_label = Variable(batch_image), Variable(batch_id_label) steps += 1 Prediction = Model(batch_image) Loss = Model.ID_Loss(Prediction, batch_id_label) #Loss = loss_criterion_Angular(Prediction, batch_id_label) Loss.backward() optimizer.step() if args.dynamic_lr == True: scheduler.step(Loss) log_learning(epoch, steps, 'ResNet50_Model', args.lr, Loss.item(), args) writer.add_scalar('Train/Train_Loss', Loss, steps) writer.add_scalar('Train/Model_Lr', optimizer.param_groups[0]['lr'], epoch) # Validation_Process(Model, epoch, writer, args) Validation_Process(Model, epoch, writer, args) if epoch % args.save_freq == 0: if not os.path.isdir(args.snapshot_dir): os.makedirs(args.snapshot_dir) save_path = os.path.join(args.snapshot_dir, 'epoch{}.pt'.format(epoch)) torch.save(Model.state_dict(), save_path) save_checkpoint({ 'epoch': epoch + 1, 'Model': Model.state_dict(), 'optimizer': optimizer.state_dict(), }, save_dir=os.path.join(args.snapshot_dir, 'epoch{}'.format(epoch))) # export scalar data to JSON for external processing writer.export_scalars_to_json("./all_scalars.json") writer.close()
def Train(Model, args): #Define num of classes Nd = args.Nd beta1_Adam = args.beta1 beta2_Adam = args.beta2 #Define gpu mode if args.cuda: Model.cuda() #choose your optimizer optimizer = optim.Adam(Model.parameters(), lr=args.lr, betas=(beta1_Adam, beta2_Adam)) if args.resume: checkpoint = torch.load(args.resume) optimizer.load_state_dict(checkpoint['optimizer']) Model.train() loss_criterion = nn.CrossEntropyLoss().cuda() steps = 0 CUDNN.benchmark = True for epoch in range(args.start_epoch, args.epochs + 1): # Every args.lr_step, changes learning rate by multipling args.lr_decay if args.step_learning: adjust_learning_rate(optimizer, epoch, args) # Load augmented data transformed_dataset = FaceIdPoseDataset(args.train_csv_file, args.data_place, transform=transforms.Compose([ Resize((256, 256)), RandomCrop((224, 224)) ])) dataloader = DataLoader(transformed_dataset, batch_size=args.Train_Batch, shuffle=True) for i, batch_data in enumerate(dataloader): # backward() function accumulates gradients, however we don't want to mix up gradients between minibatches Model.zero_grad() batch_image = torch.FloatTensor(batch_data[0].float()) batch_id_label = batch_data[2] if args.cuda: batch_image, batch_id_label = batch_image.cuda( ), batch_id_label.cuda() batch_image, batch_id_label = Variable(batch_image), Variable( batch_id_label) steps += 1 Prediction = Model(batch_image) Loss = loss_criterion(Prediction[:, :Nd], batch_id_label) Loss.backward() optimizer.step() log_learning(epoch, steps, 'VGG16_Model', args.lr, Loss.data[0], args) writer.add_scalar('Train/Train_Loss', Loss, steps) Validation_Process(Model, epoch, writer, args) if epoch % args.save_freq == 0: if not os.path.isdir(args.snapshot_dir): os.makedirs(args.snapshot_dir) save_path = os.path.join(args.snapshot_dir, 'epoch{}.pt'.format(epoch)) torch.save(Model.state_dict(), save_path) save_checkpoint( { 'epoch': epoch + 1, 'Model': Model.state_dict(), 'optimizer': optimizer.state_dict(), }, save_dir=os.path.join(args.snapshot_dir, 'epoch{}'.format(epoch))) # export scalar data to JSON for external processing writer.export_scalars_to_json("./all_scalars.json") writer.close()