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)
# 最後のエポックの学習前に生成した画像を１枚保存（学習の確認用）
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)
# 最後のエポックの学習前に生成した画像を１枚保存（学習の確認用）
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()
