def build_model(self):
"""Create a generator and a discriminator."""
if self.dataset in ['CelebA', 'RaFD']:
self.G = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim, self.d_repeat_num)
elif self.dataset in ['Both']:
self.G = Generator(self.g_conv_dim, self.c_dim+self.c2_dim+2, self.g_repeat_num) # 2 for mask vector.
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim+self.c2_dim, self.d_repeat_num)
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.G.to(self.device)
self.D.to(self.device)
def build_model(self):
# Define a generator and a discriminator
if self.dataset == 'Both':
self.G = Generator(self.g_conv_dim, self.c_dim+self.c2_dim+2, self.g_repeat_num) # 2 for mask vector
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim+self.c2_dim, self.d_repeat_num)
else:
self.G = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim, self.d_repeat_num)
# Optimizers
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])
# Print networks
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
if torch.cuda.is_available():
self.G.cuda()
self.D.cuda()
def train(resume_path=None, jigsaw_path=None):
writer = SummaryWriter('../runs/'+hparams.exp_name)
for k in hparams.__dict__.keys():
writer.add_text(str(k), str(hparams.__dict__[k]))
train_dataset = ChestData(data_csv=hparams.train_csv, data_dir=hparams.train_dir, augment=hparams.augment,
transform=transforms.Compose([
transforms.Resize(hparams.image_shape),
transforms.ToTensor(),
transforms.Normalize((0.5027, 0.5027, 0.5027), (0.2915, 0.2915, 0.2915))
]))
validation_dataset = ChestData(data_csv=hparams.valid_csv, data_dir=hparams.valid_dir,
transform=transforms.Compose([
transforms.Resize(hparams.image_shape),
transforms.ToTensor(),
transforms.Normalize((0.5027, 0.5027, 0.5027), (0.2915, 0.2915, 0.2915))
]))
# train_sampler = WeightedRandomSampler()
train_loader = DataLoader(train_dataset, batch_size=hparams.batch_size,
shuffle=True, num_workers=2)
validation_loader = DataLoader(validation_dataset, batch_size=hparams.batch_size,
shuffle=True, num_workers=2)
print('loaded train data of length : {}'.format(len(train_dataset)))
adversarial_loss = torch.nn.BCELoss().to(hparams.gpu_device)
discriminator = Discriminator().to(hparams.gpu_device)
if hparams.cuda:
discriminator = nn.DataParallel(discriminator, device_ids=hparams.device_ids)
params_count = 0
for param in discriminator.parameters():
params_count += np.prod(param.size())
print('Model has {0} trainable parameters'.format(params_count))
if not hparams.pretrained:
# discriminator.apply(weights_init_normal)
pass
if jigsaw_path:
jigsaw = Jigsaw().to(hparams.gpu_device)
if hparams.cuda:
jigsaw = nn.DataParallel(jigsaw, device_ids=hparams.device_ids)
checkpoints = torch.load(jigsaw_path, map_location=hparams.gpu_device)
jigsaw.load_state_dict(checkpoints['discriminator_state_dict'])
discriminator.module.model.features = jigsaw.module.feature.features
print('loaded pretrained feature extractor from {} ..'.format(jigsaw_path))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=hparams.learning_rate, betas=(0.9, 0.999))
scheduler_D = ReduceLROnPlateau(optimizer_D, mode='min', factor=0.1, patience=1, verbose=True, cooldown=0)
Tensor = torch.cuda.FloatTensor if hparams.cuda else torch.FloatTensor
def validation(discriminator, send_stats=False, epoch=0):
print('Validating model on {0} examples. '.format(len(validation_dataset)))
discriminator_ = discriminator.eval()
with torch.no_grad():
pred_logits_list = []
labels_list = []
for (img, labels, imgs_names) in tqdm(validation_loader):
img = Variable(img.float(), requires_grad=False)
labels = Variable(labels.float(), requires_grad=False)
img_ = img.to(hparams.gpu_device)
labels = labels.to(hparams.gpu_device)
pred_logits = discriminator_(img_)
pred_logits_list.append(pred_logits)
labels_list.append(labels)
pred_logits = torch.cat(pred_logits_list, dim=0)
labels = torch.cat(labels_list, dim=0)
val_loss = adversarial_loss(pred_logits, labels)
return accuracy_metrics(labels.long(), pred_logits), val_loss
print('Starting training.. (log saved in:{})'.format(hparams.exp_name))
start_time = time.time()
best_valid_auc = 0
# print(model)
for epoch in range(hparams.num_epochs):
for batch, (imgs, labels, imgs_name) in enumerate(tqdm(train_loader)):
imgs = Variable(imgs.float(), requires_grad=False)
labels = Variable(labels.float(), requires_grad=False)
imgs_ = imgs.to(hparams.gpu_device)
labels = labels.to(hparams.gpu_device)
# ---------------------
# Train Discriminator
# ---------------------
optimizer_D.zero_grad()
pred_logits = discriminator(imgs_)
d_loss = adversarial_loss(pred_logits, labels)
d_loss.backward()
optimizer_D.step()
writer.add_scalar('d_loss', d_loss.item(), global_step=batch+epoch*len(train_loader))
pred_labels = (pred_logits >= hparams.thresh)
pred_labels = pred_labels.float()
# if batch % hparams.print_interval == 0:
# auc, f1, acc, _, _ = accuracy_metrics(pred_labels, labels.long(), pred_logits)
# print('[Epoch - {0:.1f}, batch - {1:.3f}, d_loss - {2:.6f}, acc - {3:.4f}, f1 - {4:.5f}, auc - {5:.4f}]'.\
# format(1.0*epoch, 100.0*batch/len(train_loader), d_loss.item(), acc['avg'], f1[hparams.avg_mode], auc[hparams.avg_mode]))
(val_auc, val_f1, val_acc, val_conf_mat, best_thresh), val_loss = validation(discriminator, epoch=epoch)
for lbl in range(hparams.num_classes):
fig = plot_cf(val_conf_mat[lbl])
writer.add_figure('val_conf_{}'.format(hparams.id_to_class[lbl]), fig, global_step=epoch)
plt.close(fig)
writer.add_scalar('val_f1_{}'.format(hparams.id_to_class[lbl]), val_f1[lbl], global_step=epoch)
writer.add_scalar('val_auc_{}'.format(hparams.id_to_class[lbl]), val_auc[lbl], global_step=epoch)
writer.add_scalar('val_acc_{}'.format(hparams.id_to_class[lbl]), val_acc[lbl], global_step=epoch)
writer.add_scalar('val_f1_{}'.format('micro'), val_f1['micro'], global_step=epoch)
writer.add_scalar('val_auc_{}'.format('micro'), val_auc['micro'], global_step=epoch)
writer.add_scalar('val_f1_{}'.format('macro'), val_f1['macro'], global_step=epoch)
writer.add_scalar('val_auc_{}'.format('macro'), val_auc['macro'], global_step=epoch)
writer.add_scalar('val_loss', val_loss, global_step=epoch)
writer.add_scalar('val_f1', val_f1[hparams.avg_mode], global_step=epoch)
writer.add_scalar('val_auc', val_auc[hparams.avg_mode], global_step=epoch)
writer.add_scalar('val_acc', val_acc['avg'], global_step=epoch)
scheduler_D.step(val_loss)
writer.add_scalar('learning_rate', optimizer_D.param_groups[0]['lr'], global_step=epoch)
torch.save({
'epoch': epoch,
'discriminator_state_dict': discriminator.state_dict(),
'optimizer_D_state_dict': optimizer_D.state_dict(),
}, hparams.model+'.'+str(epoch))
if best_valid_auc <= val_auc[hparams.avg_mode]:
best_valid_auc = val_auc[hparams.avg_mode]
for lbl in range(hparams.num_classes):
fig = plot_cf(val_conf_mat[lbl])
writer.add_figure('best_val_conf_{}'.format(hparams.id_to_class[lbl]), fig, global_step=epoch)
plt.close(fig)
torch.save({
'epoch': epoch,
'discriminator_state_dict': discriminator.state_dict(),
'optimizer_D_state_dict': optimizer_D.state_dict(),
}, hparams.model+'.best')
print('best model on validation set saved.')
print('[Epoch - {0:.1f} ---> val_auc - {1:.4f}, current_lr - {2:.6f}, val_loss - {3:.4f}, best_val_auc - {4:.4f}, val_acc - {5:.4f}, val_f1 - {6:.4f}] - time - {7:.1f}'\
.format(1.0*epoch, val_auc[hparams.avg_mode], optimizer_D.param_groups[0]['lr'], val_loss, best_valid_auc, val_acc['avg'], val_f1[hparams.avg_mode], time.time()-start_time))
start_time = time.time()
class Solver(object):
def __init__(self, celebA_loader, rafd_loader, config):
# Data loader
self.celebA_loader = celebA_loader
self.rafd_loader = rafd_loader
# Model hyper-parameters
self.c_dim = config.c_dim
self.c2_dim = config.c2_dim
self.image_size = config.image_size
self.g_conv_dim = config.g_conv_dim
self.d_conv_dim = config.d_conv_dim
self.g_repeat_num = config.g_repeat_num
self.d_repeat_num = config.d_repeat_num
self.d_train_repeat = config.d_train_repeat
# Hyper-parameteres
self.lambda_cls = config.lambda_cls
self.lambda_rec = config.lambda_rec
self.lambda_gp = config.lambda_gp
self.g_lr = config.g_lr
self.d_lr = config.d_lr
self.beta1 = config.beta1
self.beta2 = config.beta2
# Training settings
self.dataset = config.dataset
self.num_epochs = config.num_epochs
self.num_epochs_decay = config.num_epochs_decay
self.num_iters = config.num_iters
self.num_iters_decay = config.num_iters_decay
self.batch_size = config.batch_size
self.use_tensorboard = config.use_tensorboard
self.pretrained_model = config.pretrained_model
# Test settings
self.test_model = config.test_model
# Path
self.log_path = config.log_path
self.sample_path = config.sample_path
self.model_save_path = config.model_save_path
self.result_path = config.result_path
# Step size
self.log_step = config.log_step
self.sample_step = config.sample_step
self.model_save_step = config.model_save_step
# Build tensorboard if use
self.build_model()
if self.use_tensorboard:
self.build_tensorboard()
# Start with trained model
if self.pretrained_model:
self.load_pretrained_model()
def build_model(self):
# Define a generator and a discriminator
if self.dataset == 'Both':
self.G = Generator(self.g_conv_dim, self.c_dim+self.c2_dim+2, self.g_repeat_num) # 2 for mask vector
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim+self.c2_dim, self.d_repeat_num)
else:
self.G = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim, self.d_repeat_num)
# Optimizers
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])
# Print networks
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
if torch.cuda.is_available():
self.G.cuda()
self.D.cuda()
def print_network(self, model, name):
num_params = 0
for p in model.parameters():
num_params += p.numel()
print(name)
print(model)
print("The number of parameters: {}".format(num_params))
def load_pretrained_model(self):
self.G.load_state_dict(torch.load(os.path.join(
self.model_save_path, '{}_G.pth'.format(self.pretrained_model))))
self.D.load_state_dict(torch.load(os.path.join(
self.model_save_path, '{}_D.pth'.format(self.pretrained_model))))
print('loaded trained models (step: {})..!'.format(self.pretrained_model))
def build_tensorboard(self):
from logger import Logger
self.logger = Logger(self.log_path)
def update_lr(self, g_lr, d_lr):
for param_group in self.g_optimizer.param_groups:
param_group['lr'] = g_lr
for param_group in self.d_optimizer.param_groups:
param_group['lr'] = d_lr
def reset_grad(self):
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
def to_var(self, x, volatile=False):
if torch.cuda.is_available():
x = x.cuda()
return Variable(x, volatile=volatile)
def denorm(self, x):
out = (x + 1) / 2
return out.clamp_(0, 1)
def threshold(self, x):
x = x.clone()
x[x >= 0.5] = 1
x[x < 0.5] = 0
return x
def compute_accuracy(self, x, y, dataset):
if dataset == 'CelebA':
x = F.sigmoid(x)
predicted = self.threshold(x)
correct = (predicted == y).float()
accuracy = torch.mean(correct, dim=0) * 100.0
else:
_, predicted = torch.max(x, dim=1)
correct = (predicted == y).float()
accuracy = torch.mean(correct) * 100.0
return accuracy
def one_hot(self, labels, dim):
"""Convert label indices to one-hot vector"""
batch_size = labels.size(0)
out = torch.zeros(batch_size, dim)
out[np.arange(batch_size), labels.long()] = 1
return out
def make_celeb_labels(self, real_c):
"""Generate domain labels for CelebA for debugging/testing.
if dataset == 'CelebA':
return single and multiple attribute changes
elif dataset == 'Both':
return single attribute changes
"""
y = [torch.FloatTensor([1, 0, 0]), # black hair
torch.FloatTensor([0, 1, 0]), # blond hair
torch.FloatTensor([0, 0, 1])] # brown hair
fixed_c_list = []
# single attribute transfer
for i in range(self.c_dim):
fixed_c = real_c.clone()
for c in fixed_c:
if i < 3:
c[:3] = y[i]
else:
c[i] = 0 if c[i] == 1 else 1 # opposite value
fixed_c_list.append(self.to_var(fixed_c, volatile=True))
# multi-attribute transfer (H+G, H+A, G+A, H+G+A)
if self.dataset == 'CelebA':
for i in range(4):
fixed_c = real_c.clone()
for c in fixed_c:
if i in [0, 1, 3]: # Hair color to brown
c[:3] = y[2]
if i in [0, 2, 3]: # Gender
c[3] = 0 if c[3] == 1 else 1
if i in [1, 2, 3]: # Aged
c[4] = 0 if c[4] == 1 else 1
fixed_c_list.append(self.to_var(fixed_c, volatile=True))
return fixed_c_list
def train(self):
"""Train StarGAN within a single dataset."""
# Set dataloader
if self.dataset == 'CelebA':
self.data_loader = self.celebA_loader
else:
self.data_loader = self.rafd_loader
# The number of iterations per epoch
iters_per_epoch = len(self.data_loader)
fixed_x = []
real_c = []
for i, (images, labels) in enumerate(self.data_loader):
fixed_x.append(images)
real_c.append(labels)
if i == 3:
break
# Fixed inputs and target domain labels for debugging
fixed_x = torch.cat(fixed_x, dim=0)
fixed_x = self.to_var(fixed_x, volatile=True)
real_c = torch.cat(real_c, dim=0)
if self.dataset == 'CelebA':
fixed_c_list = self.make_celeb_labels(real_c)
elif self.dataset == 'RaFD':
fixed_c_list = []
for i in range(self.c_dim):
fixed_c = self.one_hot(torch.ones(fixed_x.size(0)) * i, self.c_dim)
fixed_c_list.append(self.to_var(fixed_c, volatile=True))
# lr cache for decaying
g_lr = self.g_lr
d_lr = self.d_lr
# Start with trained model if exists
if self.pretrained_model:
start = int(self.pretrained_model.split('_')[0])
else:
start = 0
# Start training
start_time = time.time()
for e in range(start, self.num_epochs):
for i, (real_x, real_label) in enumerate(self.data_loader):
# Generat fake labels randomly (target domain labels)
rand_idx = torch.randperm(real_label.size(0))
fake_label = real_label[rand_idx]
if self.dataset == 'CelebA':
real_c = real_label.clone()
fake_c = fake_label.clone()
else:
real_c = self.one_hot(real_label, self.c_dim)
fake_c = self.one_hot(fake_label, self.c_dim)
# Convert tensor to variable
real_x = self.to_var(real_x)
real_c = self.to_var(real_c) # input for the generator
fake_c = self.to_var(fake_c)
real_label = self.to_var(real_label) # this is same as real_c if dataset == 'CelebA'
fake_label = self.to_var(fake_label)
# ================== Train D ================== #
# Compute loss with real images
out_src, out_cls = self.D(real_x)
d_loss_real = - torch.mean(out_src)
if self.dataset == 'CelebA':
d_loss_cls = F.binary_cross_entropy_with_logits(
out_cls, real_label, size_average=False) / real_x.size(0)
else:
d_loss_cls = F.cross_entropy(out_cls, real_label)
# Compute classification accuracy of the discriminator
if (i+1) % self.log_step == 0:
accuracies = self.compute_accuracy(out_cls, real_label, self.dataset)
log = ["{:.2f}".format(acc) for acc in accuracies.data.cpu().numpy()]
if self.dataset == 'CelebA':
print('Classification Acc (Black/Blond/Brown/Gender/Aged): ', end='')
else:
print('Classification Acc (8 emotional expressions): ', end='')
print(log)
# Compute loss with fake images
fake_x = self.G(real_x, fake_c)
fake_x = Variable(fake_x.data)
out_src, out_cls = self.D(fake_x)
d_loss_fake = torch.mean(out_src)
# Backward + Optimize
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Compute gradient penalty
alpha = torch.rand(real_x.size(0), 1, 1, 1).cuda().expand_as(real_x)
interpolated = Variable(alpha * real_x.data + (1 - alpha) * fake_x.data, requires_grad=True)
out, out_cls = self.D(interpolated)
grad = torch.autograd.grad(outputs=out,
inputs=interpolated,
grad_outputs=torch.ones(out.size()).cuda(),
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
grad = grad.view(grad.size(0), -1)
grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))
d_loss_gp = torch.mean((grad_l2norm - 1)**2)
# Backward + Optimize
d_loss = self.lambda_gp * d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Logging
loss = {}
loss['D/loss_real'] = d_loss_real.data[0]
loss['D/loss_fake'] = d_loss_fake.data[0]
loss['D/loss_cls'] = d_loss_cls.data[0]
loss['D/loss_gp'] = d_loss_gp.data[0]
# ================== Train G ================== #
if (i+1) % self.d_train_repeat == 0:
# Original-to-target and target-to-original domain
fake_x = self.G(real_x, fake_c)
rec_x = self.G(fake_x, real_c)
# Compute losses
out_src, out_cls = self.D(fake_x)
g_loss_fake = - torch.mean(out_src)
g_loss_rec = torch.mean(torch.abs(real_x - rec_x))
if self.dataset == 'CelebA':
g_loss_cls = F.binary_cross_entropy_with_logits(
out_cls, fake_label, size_average=False) / fake_x.size(0)
else:
g_loss_cls = F.cross_entropy(out_cls, fake_label)
# Backward + Optimize
g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + self.lambda_cls * g_loss_cls
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging
loss['G/loss_fake'] = g_loss_fake.data[0]
loss['G/loss_rec'] = g_loss_rec.data[0]
loss['G/loss_cls'] = g_loss_cls.data[0]
# Print out log info
if (i+1) % self.log_step == 0:
elapsed = time.time() - start_time
elapsed = str(datetime.timedelta(seconds=elapsed))
log = "Elapsed [{}], Epoch [{}/{}], Iter [{}/{}]".format(
elapsed, e+1, self.num_epochs, i+1, iters_per_epoch)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, e * iters_per_epoch + i + 1)
# Translate fixed images for debugging
if (i+1) % self.sample_step == 0:
fake_image_list = [fixed_x]
for fixed_c in fixed_c_list:
fake_image_list.append(self.G(fixed_x, fixed_c))
fake_images = torch.cat(fake_image_list, dim=3)
save_image(self.denorm(fake_images.data),
os.path.join(self.sample_path, '{}_{}_fake.png'.format(e+1, i+1)),nrow=1, padding=0)
print('Translated images and saved into {}..!'.format(self.sample_path))
# Save model checkpoints
if (i+1) % self.model_save_step == 0:
torch.save(self.G.state_dict(),
os.path.join(self.model_save_path, '{}_{}_G.pth'.format(e+1, i+1)))
torch.save(self.D.state_dict(),
os.path.join(self.model_save_path, '{}_{}_D.pth'.format(e+1, i+1)))
# Decay learning rate
if (e+1) > (self.num_epochs - self.num_epochs_decay):
g_lr -= (self.g_lr / float(self.num_epochs_decay))
d_lr -= (self.d_lr / float(self.num_epochs_decay))
self.update_lr(g_lr, d_lr)
print ('Decay learning rate to g_lr: {}, d_lr: {}.'.format(g_lr, d_lr))
def train_multi(self):
"""Train StarGAN with multiple datasets.
In the code below, 1 is related to CelebA and 2 is releated to RaFD.
"""
# Fixed imagse and labels for debugging
fixed_x = []
real_c = []
for i, (images, labels) in enumerate(self.celebA_loader):
fixed_x.append(images)
real_c.append(labels)
if i == 2:
break
fixed_x = torch.cat(fixed_x, dim=0)
fixed_x = self.to_var(fixed_x, volatile=True)
real_c = torch.cat(real_c, dim=0)
fixed_c1_list = self.make_celeb_labels(real_c)
fixed_c2_list = []
for i in range(self.c2_dim):
fixed_c = self.one_hot(torch.ones(fixed_x.size(0)) * i, self.c2_dim)
fixed_c2_list.append(self.to_var(fixed_c, volatile=True))
fixed_zero1 = self.to_var(torch.zeros(fixed_x.size(0), self.c2_dim)) # zero vector when training with CelebA
fixed_mask1 = self.to_var(self.one_hot(torch.zeros(fixed_x.size(0)), 2)) # mask vector: [1, 0]
fixed_zero2 = self.to_var(torch.zeros(fixed_x.size(0), self.c_dim)) # zero vector when training with RaFD
fixed_mask2 = self.to_var(self.one_hot(torch.ones(fixed_x.size(0)), 2)) # mask vector: [0, 1]
# lr cache for decaying
g_lr = self.g_lr
d_lr = self.d_lr
# data iterator
data_iter1 = iter(self.celebA_loader)
data_iter2 = iter(self.rafd_loader)
# Start with trained model
if self.pretrained_model:
start = int(self.pretrained_model) + 1
else:
start = 0
# # Start training
start_time = time.time()
for i in range(start, self.num_iters):
# Fetch mini-batch images and labels
try:
real_x1, real_label1 = next(data_iter1)
except:
data_iter1 = iter(self.celebA_loader)
real_x1, real_label1 = next(data_iter1)
try:
real_x2, real_label2 = next(data_iter2)
except:
data_iter2 = iter(self.rafd_loader)
real_x2, real_label2 = next(data_iter2)
# Generate fake labels randomly (target domain labels)
rand_idx = torch.randperm(real_label1.size(0))
fake_label1 = real_label1[rand_idx]
rand_idx = torch.randperm(real_label2.size(0))
fake_label2 = real_label2[rand_idx]
real_c1 = real_label1.clone()
fake_c1 = fake_label1.clone()
zero1 = torch.zeros(real_x1.size(0), self.c2_dim)
mask1 = self.one_hot(torch.zeros(real_x1.size(0)), 2)
real_c2 = self.one_hot(real_label2, self.c2_dim)
fake_c2 = self.one_hot(fake_label2, self.c2_dim)
zero2 = torch.zeros(real_x2.size(0), self.c_dim)
mask2 = self.one_hot(torch.ones(real_x2.size(0)), 2)
# Convert tensor to variable
real_x1 = self.to_var(real_x1)
real_c1 = self.to_var(real_c1)
fake_c1 = self.to_var(fake_c1)
mask1 = self.to_var(mask1)
zero1 = self.to_var(zero1)
real_x2 = self.to_var(real_x2)
real_c2 = self.to_var(real_c2)
fake_c2 = self.to_var(fake_c2)
mask2 = self.to_var(mask2)
zero2 = self.to_var(zero2)
real_label1 = self.to_var(real_label1)
fake_label1 = self.to_var(fake_label1)
real_label2 = self.to_var(real_label2)
fake_label2 = self.to_var(fake_label2)
# ================== Train D ================== #
# Real images (CelebA)
out_real, out_cls = self.D(real_x1)
out_cls1 = out_cls[:, :self.c_dim] # celebA part
d_loss_real = - torch.mean(out_real)
d_loss_cls = F.binary_cross_entropy_with_logits(out_cls1, real_label1, size_average=False) / real_x1.size(0)
# Real images (RaFD)
out_real, out_cls = self.D(real_x2)
out_cls2 = out_cls[:, self.c_dim:] # rafd part
d_loss_real += - torch.mean(out_real)
d_loss_cls += F.cross_entropy(out_cls2, real_label2)
# Compute classification accuracy of the discriminator
if (i+1) % self.log_step == 0:
accuracies = self.compute_accuracy(out_cls1, real_label1, 'CelebA')
log = ["{:.2f}".format(acc) for acc in accuracies.data.cpu().numpy()]
print('Classification Acc (Black/Blond/Brown/Gender/Aged): ', end='')
print(log)
accuracies = self.compute_accuracy(out_cls2, real_label2, 'RaFD')
log = ["{:.2f}".format(acc) for acc in accuracies.data.cpu().numpy()]
print('Classification Acc (8 emotional expressions): ', end='')
print(log)
# Fake images (CelebA)
fake_c = torch.cat([fake_c1, zero1, mask1], dim=1)
fake_x1 = self.G(real_x1, fake_c)
fake_x1 = Variable(fake_x1.data)
out_fake, _ = self.D(fake_x1)
d_loss_fake = torch.mean(out_fake)
# Fake images (RaFD)
fake_c = torch.cat([zero2, fake_c2, mask2], dim=1)
fake_x2 = self.G(real_x2, fake_c)
out_fake, _ = self.D(fake_x2)
d_loss_fake += torch.mean(out_fake)
# Backward + Optimize
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Compute gradient penalty
if (i+1) % 2 == 0:
real_x = real_x1
fake_x = fake_x1
else:
real_x = real_x2
fake_x = fake_x2
alpha = torch.rand(real_x.size(0), 1, 1, 1).cuda().expand_as(real_x)
interpolated = Variable(alpha * real_x.data + (1 - alpha) * fake_x.data, requires_grad=True)
out, out_cls = self.D(interpolated)
if (i+1) % 2 == 0:
out_cls = out_cls[:, :self.c_dim] # CelebA
else:
out_cls = out_cls[:, self.c_dim:] # RaFD
grad = torch.autograd.grad(outputs=out,
inputs=interpolated,
grad_outputs=torch.ones(out.size()).cuda(),
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
grad = grad.view(grad.size(0), -1)
grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))
d_loss_gp = torch.mean((grad_l2norm - 1)**2)
# Backward + Optimize
d_loss = self.lambda_gp * d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Logging
loss = {}
loss['D/loss_real'] = d_loss_real.data[0]
loss['D/loss_fake'] = d_loss_fake.data[0]
loss['D/loss_cls'] = d_loss_cls.data[0]
loss['D/loss_gp'] = d_loss_gp.data[0]
# ================== Train G ================== #
if (i+1) % self.d_train_repeat == 0:
# Original-to-target and target-to-original domain (CelebA)
fake_c = torch.cat([fake_c1, zero1, mask1], dim=1)
real_c = torch.cat([real_c1, zero1, mask1], dim=1)
fake_x1 = self.G(real_x1, fake_c)
rec_x1 = self.G(fake_x1, real_c)
# Compute losses
out, out_cls = self.D(fake_x1)
out_cls1 = out_cls[:, :self.c_dim]
g_loss_fake = - torch.mean(out)
g_loss_rec = torch.mean(torch.abs(real_x1 - rec_x1))
g_loss_cls = F.binary_cross_entropy_with_logits(out_cls1, fake_label1, size_average=False) / fake_x1.size(0)
# Original-to-target and target-to-original domain (RaFD)
fake_c = torch.cat([zero2, fake_c2, mask2], dim=1)
real_c = torch.cat([zero2, real_c2, mask2], dim=1)
fake_x2 = self.G(real_x2, fake_c)
rec_x2 = self.G(fake_x2, real_c)
# Compute losses
out, out_cls = self.D(fake_x2)
out_cls2 = out_cls[:, self.c_dim:]
g_loss_fake += - torch.mean(out)
g_loss_rec += torch.mean(torch.abs(real_x2 - rec_x2))
g_loss_cls += F.cross_entropy(out_cls2, fake_label2)
# Backward + Optimize
g_loss = g_loss_fake + self.lambda_cls * g_loss_cls + self.lambda_rec * g_loss_rec
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging
loss['G/loss_fake'] = g_loss_fake.data[0]
loss['G/loss_cls'] = g_loss_cls.data[0]
loss['G/loss_rec'] = g_loss_rec.data[0]
# Print out log info
if (i+1) % self.log_step == 0:
elapsed = time.time() - start_time
elapsed = str(datetime.timedelta(seconds=elapsed))
log = "Elapsed [{}], Iter [{}/{}]".format(
elapsed, i+1, self.num_iters)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, i+1)
# Translate the images (debugging)
if (i+1) % self.sample_step == 0:
fake_image_list = [fixed_x]
# Changing hair color, gender, and age
for j in range(self.c_dim):
fake_c = torch.cat([fixed_c1_list[j], fixed_zero1, fixed_mask1], dim=1)
fake_image_list.append(self.G(fixed_x, fake_c))
# Changing emotional expressions
for j in range(self.c2_dim):
fake_c = torch.cat([fixed_zero2, fixed_c2_list[j], fixed_mask2], dim=1)
fake_image_list.append(self.G(fixed_x, fake_c))
fake = torch.cat(fake_image_list, dim=3)
# Save the translated images
save_image(self.denorm(fake.data),
os.path.join(self.sample_path, '{}_fake.png'.format(i+1)), nrow=1, padding=0)
# Save model checkpoints
if (i+1) % self.model_save_step == 0:
torch.save(self.G.state_dict(),
os.path.join(self.model_save_path, '{}_G.pth'.format(i+1)))
torch.save(self.D.state_dict(),
os.path.join(self.model_save_path, '{}_D.pth'.format(i+1)))
# Decay learning rate
decay_step = 1000
if (i+1) > (self.num_iters - self.num_iters_decay) and (i+1) % decay_step==0:
g_lr -= (self.g_lr / float(self.num_iters_decay) * decay_step)
d_lr -= (self.d_lr / float(self.num_iters_decay) * decay_step)
self.update_lr(g_lr, d_lr)
print ('Decay learning rate to g_lr: {}, d_lr: {}.'.format(g_lr, d_lr))
def test(self):
"""Facial attribute transfer on CelebA or facial expression synthesis on RaFD."""
# Load trained parameters
G_path = os.path.join(self.model_save_path, '{}_G.pth'.format(self.test_model))
self.G.load_state_dict(torch.load(G_path))
self.G.eval()
if self.dataset == 'CelebA':
data_loader = self.celebA_loader
else:
data_loader = self.rafd_loader
for i, (real_x, org_c) in enumerate(data_loader):
real_x = self.to_var(real_x, volatile=True)
if self.dataset == 'CelebA':
target_c_list = self.make_celeb_labels(org_c)
else:
target_c_list = []
for j in range(self.c_dim):
target_c = self.one_hot(torch.ones(real_x.size(0)) * j, self.c_dim)
target_c_list.append(self.to_var(target_c, volatile=True))
# Start translations
fake_image_list = [real_x]
for target_c in target_c_list:
fake_image_list.append(self.G(real_x, target_c))
fake_images = torch.cat(fake_image_list, dim=3)
save_path = os.path.join(self.result_path, '{}_fake.png'.format(i+1))
save_image(self.denorm(fake_images.data), save_path, nrow=1, padding=0)
print('Translated test images and saved into "{}"..!'.format(save_path))
def test_multi(self):
"""Facial attribute transfer and expression synthesis on CelebA."""
# Load trained parameters
G_path = os.path.join(self.model_save_path, '{}_G.pth'.format(self.test_model))
self.G.load_state_dict(torch.load(G_path))
self.G.eval()
for i, (real_x, org_c) in enumerate(self.celebA_loader):
# Prepare input images and target domain labels
real_x = self.to_var(real_x, volatile=True)
target_c1_list = self.make_celeb_labels(org_c)
target_c2_list = []
for j in range(self.c2_dim):
target_c = self.one_hot(torch.ones(real_x.size(0)) * j, self.c2_dim)
target_c2_list.append(self.to_var(target_c, volatile=True))
# Zero vectors and mask vectors
zero1 = self.to_var(torch.zeros(real_x.size(0), self.c2_dim)) # zero vector for rafd expressions
mask1 = self.to_var(self.one_hot(torch.zeros(real_x.size(0)), 2)) # mask vector: [1, 0]
zero2 = self.to_var(torch.zeros(real_x.size(0), self.c_dim)) # zero vector for celebA attributes
mask2 = self.to_var(self.one_hot(torch.ones(real_x.size(0)), 2)) # mask vector: [0, 1]
# Changing hair color, gender, and age
fake_image_list = [real_x]
for j in range(self.c_dim):
target_c = torch.cat([target_c1_list[j], zero1, mask1], dim=1)
fake_image_list.append(self.G(real_x, target_c))
# Changing emotional expressions
for j in range(self.c2_dim):
target_c = torch.cat([zero2, target_c2_list[j], mask2], dim=1)
fake_image_list.append(self.G(real_x, target_c))
fake_images = torch.cat(fake_image_list, dim=3)
# Save the translated images
save_path = os.path.join(self.result_path, '{}_fake.png'.format(i+1))
save_image(self.denorm(fake_images.data), save_path, nrow=1, padding=0)
print('Translated test images and saved into "{}"..!'.format(save_path))
def main():
env = gym.make(args.env_name)
env.seed(args.seed)
torch.manual_seed(args.seed)
num_inputs = env.observation_space.shape[0]
num_actions = env.action_space.shape[0]
running_state = ZFilter((num_inputs,), clip=5)
print('state size:', num_inputs)
print('action size:', num_actions)
actor = Actor(num_inputs, num_actions, args)
critic = Critic(num_inputs, args)
discrim = Discriminator(num_inputs + num_actions, args)
actor_optim = optim.Adam(actor.parameters(), lr=args.learning_rate)
critic_optim = optim.Adam(critic.parameters(), lr=args.learning_rate,
weight_decay=args.l2_rate)
discrim_optim = optim.Adam(discrim.parameters(), lr=args.learning_rate)
# load demonstrations
expert_demo, _ = pickle.load(open('./expert_demo/expert_demo.p', "rb"))
demonstrations = np.array(expert_demo)
print("demonstrations.shape", demonstrations.shape)
writer = SummaryWriter(args.logdir)
if args.load_model is not None:
saved_ckpt_path = os.path.join(os.getcwd(), 'save_model', str(args.load_model))
ckpt = torch.load(saved_ckpt_path)
actor.load_state_dict(ckpt['actor'])
critic.load_state_dict(ckpt['critic'])
discrim.load_state_dict(ckpt['discrim'])
running_state.rs.n = ckpt['z_filter_n']
running_state.rs.mean = ckpt['z_filter_m']
running_state.rs.sum_square = ckpt['z_filter_s']
print("Loaded OK ex. Zfilter N {}".format(running_state.rs.n))
episodes = 0
train_discrim_flag = True
for iter in range(args.max_iter_num):
actor.eval(), critic.eval()
memory = deque()
steps = 0
scores = []
while steps < args.total_sample_size:
state = env.reset()
score = 0
state = running_state(state)
for _ in range(10000):
if args.render:
env.render()
steps += 1
mu, std = actor(torch.Tensor(state).unsqueeze(0))
action = get_action(mu, std)[0]
next_state, reward, done, _ = env.step(action)
irl_reward = get_reward(discrim, state, action)
if done:
mask = 0
else:
mask = 1
memory.append([state, action, irl_reward, mask])
next_state = running_state(next_state)
state = next_state
score += reward
if done:
break
episodes += 1
scores.append(score)
score_avg = np.mean(scores)
print('{}:: {} episode score is {:.2f}'.format(iter, episodes, score_avg))
writer.add_scalar('log/score', float(score_avg), iter)
actor.train(), critic.train(), discrim.train()
if train_discrim_flag:
expert_acc, learner_acc = train_discrim(discrim, memory, discrim_optim, demonstrations, args)
print("Expert: %.2f%% | Learner: %.2f%%" % (expert_acc * 100, learner_acc * 100))
if expert_acc > args.suspend_accu_exp and learner_acc > args.suspend_accu_gen:
train_discrim_flag = False
train_actor_critic(actor, critic, memory, actor_optim, critic_optim, args)
if iter % 100:
score_avg = int(score_avg)
model_path = os.path.join(os.getcwd(),'save_model')
if not os.path.isdir(model_path):
os.makedirs(model_path)
ckpt_path = os.path.join(model_path, 'ckpt_'+ str(score_avg)+'.pth.tar')
save_checkpoint({
'actor': actor.state_dict(),
'critic': critic.state_dict(),
'discrim': discrim.state_dict(),
'z_filter_n':running_state.rs.n,
'z_filter_m': running_state.rs.mean,
'z_filter_s': running_state.rs.sum_square,
'args': args,
'score': score_avg
}, filename=ckpt_path)
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method="env://")
synchronize()
args.latent = 512
args.n_mlp = 8
args.start_iter = 0
generator = Generator(
args.size,
args.latent,
args.n_mlp,
channel_multiplier=args.channel_multiplier).to(device)
discriminator = Discriminator(
args.size, channel_multiplier=args.channel_multiplier).to(device)
g_ema = Generator(args.size,
args.latent,
args.n_mlp,
channel_multiplier=args.channel_multiplier).to(device)
g_ema.eval()
accumulate(g_ema, generator, 0)
g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)
g_optim = optim.Adam(
generator.parameters(),
lr=args.lr * g_reg_ratio,
betas=(0**g_reg_ratio, 0.99**g_reg_ratio),
)
class Solver(object):
"""Solver for training and testing StarGAN."""
def __init__(self, celeba_loader, config):
"""Initialize configurations."""
# Data loader.
self.celeba_loader = celeba_loader
#self.rafd_loader = rafd_loader
# Model configurations.
self.c_dim = config.c_dim
#self.c2_dim = config.c2_dim
self.image_size = config.image_size
self.g_conv_dim = config.g_conv_dim
self.d_conv_dim = config.d_conv_dim
self.g_repeat_num = config.g_repeat_num
self.d_repeat_num = config.d_repeat_num
self.lambda_cls = config.lambda_cls
self.lambda_rec = config.lambda_rec
self.lambda_gp = config.lambda_gp
# Training configurations.
self.dataset = 'CelebA'
self.batch_size = config.batch_size
self.num_iters = config.num_iters
self.num_iters_decay = config.num_iters_decay
self.g_lr = config.g_lr
self.d_lr = config.d_lr
self.n_critic = config.n_critic
self.beta1 = config.beta1
self.beta2 = config.beta2
self.resume_iters = config.resume_iters
self.selected_attrs = config.selected_attrs
# Test configurations.
self.test_iters = config.test_iters
# Miscellaneous.
self.use_tensorboard = config.use_tensorboard
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
# Directories.
self.log_dir = config.log_dir
self.sample_dir = config.sample_dir
self.model_save_dir = config.model_save_dir
self.result_dir = config.result_dir
# Step size.
self.log_step = config.log_step
self.sample_step = config.sample_step
self.model_save_step = config.model_save_step
self.lr_update_step = config.lr_update_step
# Build the model and tensorboard.
self.build_model()
if self.use_tensorboard:
self.build_tensorboard()
def build_model(self):
"""Create a generator and a discriminator."""
self.G = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim,
self.d_repeat_num)
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr,
[self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr,
[self.beta1, self.beta2])
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.G.to(self.device)
self.D.to(self.device)
def print_network(self, model, name):
"""Print out the network information."""
num_params = 0
for p in model.parameters():
num_params += p.numel()
print(model)
print(name)
print("The number of parameters: {}".format(num_params))
def restore_model(self, resume_iters):
"""Restore the trained generator and discriminator."""
print(
'Loading the trained models from step {}...'.format(resume_iters))
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(resume_iters))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(resume_iters))
self.G.load_state_dict(
torch.load(G_path, map_location=lambda storage, loc: storage))
self.D.load_state_dict(
torch.load(D_path, map_location=lambda storage, loc: storage))
def build_tensorboard(self):
"""Build a tensorboard logger."""
from logger import Logger
self.logger = Logger(self.log_dir)
def update_lr(self, g_lr, d_lr):
"""Decay learning rates of the generator and discriminator."""
for param_group in self.g_optimizer.param_groups:
param_group['lr'] = g_lr
for param_group in self.d_optimizer.param_groups:
param_group['lr'] = d_lr
def reset_grad(self):
"""Reset the gradient buffers."""
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
def denorm(self, x):
"""Convert the range from [-1, 1] to [0, 1]."""
out = (x + 1) / 2
return out.clamp_(0, 1)
def gradient_penalty(self, y, x):
"""Compute gradient penalty: (L2_norm(dy/dx) - 1)**2."""
weight = torch.ones(y.size()).to(self.device)
dydx = torch.autograd.grad(outputs=y,
inputs=x,
grad_outputs=weight,
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
dydx = dydx.view(dydx.size(0), -1)
dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
return torch.mean((dydx_l2norm - 1)**2)
def create_labels(self,
c_org,
c_dim=5,
dataset='CelebA',
selected_attrs=None):
"""Generate target domain labels for debugging and testing."""
c_trg_list = []
i = 0
c_trg = c_org.clone()
c_trg[:, i] = (c_trg[:, i] == 0) # Reverse attribute value.
c_trg_list.append(c_trg.to(self.device))
#print('c_org',c_org) Reverses the original labels, for Male
#print('c_trg_list',c_trg_list)
return c_trg_list
def classification_loss(self, logit, target, dataset='CelebA'):
"""Compute binary or softmax cross entropy loss."""
return F.binary_cross_entropy_with_logits(
logit, target, size_average=False) / logit.size(0)
def train(self):
"""Train StarGAN within a single dataset."""
# Set data loader.
data_loader = self.celeba_loader
# Fetch fixed inputs for debugging.
data_iter = iter(data_loader)
x_fixed, c_org = next(data_iter)
x_fixed = x_fixed.to(self.device)
c_fixed_list = self.create_labels(c_org, self.c_dim, self.dataset,
self.selected_attrs)
# Learning rate cache for decaying.
g_lr = self.g_lr
d_lr = self.d_lr
# Start training from scratch or resume training.
start_iters = 0
if self.resume_iters:
start_iters = self.resume_iters
self.restore_model(self.resume_iters)
# Start training.
print('Start training...')
start_time = time.time()
for i in range(start_iters, (self.num_iters)):
# =================================================================================== #
# 1. Preprocess input data #
# =================================================================================== #
# Fetch real images and labels.
try:
x_real, label_org = next(data_iter)
#print(type(x_real),x_real.size()) # <class 'torch.Tensor'> torch.Size([16, 3, 128, 128])
#print(type(label_org),label_org.size()) # <class 'torch.Tensor'> torch.Size([16, 1])
except:
data_iter = iter(data_loader)
x_real, label_org = next(data_iter)
#print('asdfs')
# Generate target domain labels randomly.
rand_idx = torch.randperm(label_org.size(0))
label_trg = label_org[rand_idx]
#print(label_trg)
c_org = label_org.clone() #Actual labels from list_attr_celeb.txt
c_trg = label_trg.clone() #Batch size(16) generated random labels
x_real = x_real.to(self.device) # Input images.
c_org = c_org.to(self.device) # Original domain labels.
c_trg = c_trg.to(self.device) # Target domain labels.
label_org = label_org.to(
self.device) # Labels for computing classification loss.
label_trg = label_trg.to(
self.device) # Labels for computing classification loss.
# =================================================================================== #
# 2. Train the discriminator #
# =================================================================================== #
# Compute loss with real images.
out_src, out_cls = self.D(x_real)
#print(type(out_src),out_src.size()) #<class 'torch.Tensor'> torch.Size([16, 1, 2, 2])
#print(type(out_cls),out_cls.size()) # <class 'torch.Tensor'> torch.Size([16, 1])
d_loss_real = -torch.mean(out_src)
d_loss_cls = self.classification_loss(out_cls, label_org,
self.dataset)
# Compute loss with fake images.
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake.detach())
d_loss_fake = torch.mean(out_src)
# Compute loss for gradient penalty.
alpha = torch.rand(x_real.size(0), 1, 1, 1).to(self.device)
x_hat = (alpha * x_real.data +
(1 - alpha) * x_fake.data).requires_grad_(True)
out_src, _ = self.D(x_hat)
print('out_src,out_src.size()0', out_src, out_src.size())
print('x_hat,x_hat.size()', x_hat, x_hat.size())
d_loss_gp = self.gradient_penalty(out_src, x_hat)
print('d_loss_gp', d_loss_gp)
# Backward and optimize.
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls + self.lambda_gp * d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Logging.
loss = {}
loss['D/loss_real'] = d_loss_real.item()
loss['D/loss_fake'] = d_loss_fake.item()
loss['D/loss_cls'] = d_loss_cls.item()
loss['D/loss_gp'] = d_loss_gp.item()
# =================================================================================== #
# 3. Train the generator #
# =================================================================================== #
if (i + 1) % self.n_critic == 0:
# Original-to-target domain.
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake)
g_loss_fake = -torch.mean(out_src)
g_loss_cls = self.classification_loss(out_cls, label_trg,
self.dataset)
# Target-to-original domain.
x_reconst = self.G(x_fake, c_org)
g_loss_rec = torch.mean(torch.abs(x_real - x_reconst))
# Backward and optimize.
g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + self.lambda_cls * g_loss_cls
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging.
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_cls'] = g_loss_cls.item()
# =================================================================================== #
# 4. Miscellaneous #
# =================================================================================== #
# Print out training information.
if (i + 1) % self.log_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
log = "Elapsed [{}], Iteration [{}/{}]".format(
et, i + 1, self.num_iters)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, i + 1)
# Translate fixed images for debugging.
if (i + 1) % self.sample_step == 0:
with torch.no_grad():
x_fake_list = [x_fixed]
for c_fixed in c_fixed_list:
x_fake_list.append(self.G(x_fixed, c_fixed))
x_concat = torch.cat(x_fake_list, dim=3)
sample_path = os.path.join(self.sample_dir,
'{}-images.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()),
sample_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(
sample_path))
# Save model checkpoints.
if (i + 1) % self.model_save_step == 0:
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(i + 1))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(i + 1))
torch.save(self.G.state_dict(), G_path)
torch.save(self.D.state_dict(), D_path)
print('Saved model checkpoints into {}...'.format(
self.model_save_dir))
# Decay learning rates.
if (i + 1) % self.lr_update_step == 0 and (i + 1) > (
self.num_iters - self.num_iters_decay):
g_lr -= (self.g_lr / float(self.num_iters_decay))
d_lr -= (self.d_lr / float(self.num_iters_decay))
self.update_lr(g_lr, d_lr)
print('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(
g_lr, d_lr))
def test(self):
"""Translate images using StarGAN trained on a single dataset."""
# Load the trained generator.
self.restore_model(self.test_iters)
# Set data loader.
data_loader = self.celeba_loader
with torch.no_grad():
for i, (x_real, c_org) in enumerate(data_loader):
# Prepare input images and target domain labels.
x_real = x_real.to(self.device)
c_trg_list = self.create_labels(c_org, self.c_dim,
self.dataset,
self.selected_attrs)
# Translate images.
x_fake_list = [x_real]
for c_trg in c_trg_list:
x_fake_list.append(self.G(x_real, c_trg))
# Save the translated images.
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.result_dir,
'{}-images.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()),
result_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(
result_path))
state_dict = fill_statedict(state_dict, g_ema.vars, size)
g.load_state_dict(state_dict)
latent_avg = torch.from_numpy(g_ema.vars['dlatent_avg'].value().eval())
ckpt = {'g_ema': state_dict, 'latent_avg': latent_avg}
if args.gen:
g_train = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier)
g_train_state = g_train.state_dict()
g_train_state = fill_statedict(g_train_state, generator.vars, size)
ckpt['g'] = g_train_state
if args.disc:
disc = Discriminator(size, channel_multiplier=args.channel_multiplier)
d_state = disc.state_dict()
d_state = discriminator_fill_statedict(d_state, discriminator.vars, size)
ckpt['d'] = d_state
name = os.path.splitext(os.path.basename(args.path))[0]
torch.save(ckpt, name + '.pt')
torch.save(g, name + '_g.pt')
batch_size = {256: 16, 512: 9, 1024: 4}
n_sample = batch_size.get(size, 25)
g = g.to(device)
z = np.random.RandomState(0).randn(n_sample, 512).astype('float32')
transforms.Resize(IMAGE_SIZE),
transforms.ToTensor(),
transforms.Normalize(
[0.5 for _ in range(CHANNELS_IMG)], [0.5 for _ in range(CHANNELS_IMG)]
),
]
)
dataset = datasets.MNIST(root="dataset/", transform=transforms, download=True)
#comment mnist and uncomment below if you want to train on CelebA dataset
#dataset = datasets.ImageFolder(root="celeb_dataset", transform=transforms)
loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
# initialize gen and disc/critic
gen = Generator(Z_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)
critic = Discriminator(CHANNELS_IMG, FEATURES_CRITIC).to(device)
initialize_weights(gen)
initialize_weights(critic)
# initializate optimizer
opt_gen = optim.RMSprop(gen.parameters(), lr=LEARNING_RATE)
opt_critic = optim.RMSprop(critic.parameters(), lr=LEARNING_RATE)
# for tensorboard plotting
fixed_noise = torch.randn(32, Z_DIM, 1, 1).to(device)
writer_real = SummaryWriter(f"logs/real")
writer_fake = SummaryWriter(f"logs/fake")
step = 0
gen.train()
critic.train()
train_loader = torch.utils.data.DataLoader(X_train, batch_size=batch_size, shuffle=True, pin_memory=False) # better than for loop
val_loader = torch.utils.data.DataLoader(y_train, batch_size=batch_size, shuffle=False, pin_memory=False) # better than for loop
X_train,y_train,X_test,y_test = None,None,None,None
else:
train_set = TrainDatasetFromFolder('data/DIV2K_train_HR', crop_size=CROP_SIZE, upscale_factor=UPSCALE_FACTOR)
val_set = ValDatasetFromFolder('data/DIV2K_valid_HR', crop_size=CROP_SIZE, upscale_factor=UPSCALE_FACTOR)
train_loader = DataLoader(dataset=train_set, num_workers=4, batch_size=64, shuffle=True)
val_loader = DataLoader(dataset=val_set, num_workers=4, batch_size=1, shuffle=False)
#x,y = next(iter(val_loader)),next(iter(train_loader))
#print(x[0].shape,x[1].shape,x[2].shape,y[0].shape)
#netG = Generator(UPSCALE_FACTOR,in_channels,out_channels)
netG = UNet(n_channels=in_channels, n_classes=out_channels)
#print(summary(netG,(in_channels,128,128)))
print('# generator parameters:', sum(param.numel() for param in netG.parameters()))
netD = Discriminator(out_channels)
print('# discriminator parameters:', sum(param.numel() for param in netD.parameters()))
#print(summary(netD,(out_channels,256,256)))
generator_criterion = GeneratorLoss()
#print(summary(generator_criterion,(3,256,256)))
if torch.cuda.is_available():
netG.cuda()
netD.cuda()
generator_criterion.cuda()
optimizerG = optim.Adam(netG.parameters())
optimizerD = optim.Adam(netD.parameters())
results = {'d_loss': [], 'g_loss': [], 'd_score': [], 'g_score': [], 'psnr': [], 'ssim': []}
for epoch in range(1, NUM_EPOCHS + 1):
class Solver(object):
"""Solver for training and testing StarGAN."""
def __init__(self, celeba_loader, config):
"""Initialize configurations."""
# Data loader.
self.celeba_loader = celeba_loader
# Model configurations.
self.c_dim = config['c_dim']
self.image_size = config['image_size']
self.g_conv_dim = config['g_conv_dim']
self.d_conv_dim = config['d_conv_dim']
self.g_repeat_num = config['g_repeat_num']
self.d_repeat_num = config['d_repeat_num']
self.lambda_cls = config['lambda_cls']
self.lambda_rec = config['lambda_rec']
self.lambda_gp = config['lambda_gp']
# Training configurations.
self.dataset = config['dataset']
self.batch_size = config['batch_size']
self.num_iters = config['num_iters']
self.num_iters_decay = config['num_iters_decay']
self.g_lr = config['g_lr']
self.d_lr = config['d_lr']
self.n_critic = config['n_critic']
self.beta1 = config['beta1']
self.beta2 = config['beta2']
self.resume_iters = config['resume_iters']
self.selected_attrs = config['selected_attrs']
# Test configurations.
self.test_iters = config['test_iters']
# Miscellaneous.
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
# Directories.
self.log_dir = config['log_dir']
self.sample_dir = config['sample_dir']
self.model_save_dir = config['model_save_dir']
self.result_dir = config['result_dir']
# Step size.
self.log_step = config['log_step']
self.sample_step = config['sample_step']
self.model_save_step = config['model_save_step']
self.lr_update_step = config['lr_update_step']
# Build the model and tensorboard.
self.build_model()
def build_model(self):
"""Create a generator and a discriminator."""
self.G = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim,
self.d_repeat_num)
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr,
[self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr,
[self.beta1, self.beta2])
# self.print_network(self.G, 'G')
# self.print_network(self.D, 'D')
self.G.to(self.device)
self.D.to(self.device)
# def print_network(self, model, name):
# """Print out the network information."""
# num_params = 0
# for p in model.parameters():
# num_params += p.numel()
# print(model)
# print(name)
# print("The number of parameters: {}".format(num_params))
def restore_model(self, resume_iters):
"""Restore the trained generator and discriminator."""
print(
'Loading the trained models from step {}...'.format(resume_iters))
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(resume_iters))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(resume_iters))
self.G.load_state_dict(
torch.load(G_path, map_location=lambda storage, loc: storage))
self.D.load_state_dict(
torch.load(D_path, map_location=lambda storage, loc: storage))
def build_tensorboard(self):
"""Build a tensorboard logger."""
from logger import Logger
self.logger = Logger(self.log_dir)
def update_lr(self, g_lr, d_lr):
"""Decay learning rates of the generator and discriminator."""
for param_group in self.g_optimizer.param_groups:
param_group['lr'] = g_lr
for param_group in self.d_optimizer.param_groups:
param_group['lr'] = d_lr
def reset_grad(self):
"""Reset the gradient buffers."""
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
def denorm(self, x):
"""Convert the range from [-1, 1] to [0, 1]."""
out = (x + 1) / 2
return out.clamp_(0, 1)
def gradient_penalty(self, y, x):
"""Compute gradient penalty: (L2_norm(dy/dx) - 1)**2."""
weight = torch.ones(y.size()).to(self.device)
dydx = torch.autograd.grad(outputs=y,
inputs=x,
grad_outputs=weight,
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
dydx = dydx.view(dydx.size(0), -1)
dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
return torch.mean((dydx_l2norm - 1)**2)
def label2onehot(self, labels, dim):
"""Convert label indices to one-hot vectors."""
batch_size = labels.size(0)
out = torch.zeros(batch_size, dim)
out[np.arange(batch_size), labels.long()] = 1
return out
def create_labels(self,
c_org,
c_dim=6,
dataset='CelebA',
selected_attrs=None):
"""Generate target domain labels for debugging and testing."""
# Get hair color indices.
if dataset == 'CelebA':
hair_color_indices = []
for i, attr_name in enumerate(selected_attrs):
print("i :", i, "attr_name:", attr_name)
if attr_name in [
'Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Gray_Hair'
]:
hair_color_indices.append(i)
print("hair_color_indices: ", hair_color_indices)
c_trg_list = []
for i in range(c_dim):
c_trg = c_org.clone()
if i in hair_color_indices: # Set one hair color to 1 and the rest to 0.
c_trg[:, i] = 1
for j in hair_color_indices:
if j != i:
c_trg[:, j] = 0
else:
c_trg[:, i] = (c_trg[:, i] == 0) # Reverse attribute value.
c_trg_list.append(c_trg.to(self.device))
return c_trg_list
def classification_loss(self, logit, target, dataset='CelebA'):
"""Compute binary or softmax cross entropy loss."""
return F.binary_cross_entropy_with_logits(
logit, target, size_average=False) / logit.size(0)
def train(self):
"""Train StarGAN within a single dataset."""
# Set data loader.
data_loader = self.celeba_loader
# Fetch fixed inputs for debugging.
data_iter = iter(data_loader)
x_fixed, c_org = next(data_iter)
x_fixed = x_fixed.to(self.device)
c_fixed_list = self.create_labels(c_org, self.c_dim, self.dataset,
self.selected_attrs)
# Learning rate cache for decaying.
g_lr = self.g_lr
d_lr = self.d_lr
# Start training from scratch or resume training.
start_iters = 0
if self.resume_iters:
start_iters = self.resume_iters
self.restore_model(self.resume_iters)
# Start training.
print('Start training...')
start_time = time.time()
for i in range(start_iters, self.num_iters):
# =================================================================================== #
# 1. Preprocess input data #
# =================================================================================== #
# Fetch real images and labels.
try:
x_real, label_org = next(data_iter)
except:
data_iter = iter(data_loader)
x_real, label_org = next(data_iter)
# Generate target domain labels randomly.
rand_idx = torch.randperm(label_org.size(0))
label_trg = label_org[rand_idx]
c_org = label_org.clone()
c_trg = label_trg.clone()
x_real = x_real.to(self.device) # Input images.
c_org = c_org.to(self.device) # Original domain labels.
c_trg = c_trg.to(self.device) # Target domain labels.
label_org = label_org.to(
self.device) # Labels for computing classification loss.
label_trg = label_trg.to(
self.device) # Labels for computing classification loss.
# =================================================================================== #
# 2. Train the discriminator #
# =================================================================================== #
# Compute loss with real images.
out_src, out_cls = self.D(x_real)
d_loss_real = -torch.mean(out_src)
d_loss_cls = self.classification_loss(out_cls, label_org,
self.dataset)
# Compute loss with fake images.
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake.detach())
d_loss_fake = torch.mean(out_src)
# Compute loss for gradient penalty.
alpha = torch.rand(x_real.size(0), 1, 1, 1).to(self.device)
x_hat = (alpha * x_real.data +
(1 - alpha) * x_fake.data).requires_grad_(True)
out_src, _ = self.D(x_hat)
d_loss_gp = self.gradient_penalty(out_src, x_hat)
# Backward and optimize.
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls + self.lambda_gp * d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Logging.
loss = {}
loss['D/loss_real'] = d_loss_real.item()
loss['D/loss_fake'] = d_loss_fake.item()
loss['D/loss_cls'] = d_loss_cls.item()
loss['D/loss_gp'] = d_loss_gp.item()
# =================================================================================== #
# 3. Train the generator #
# =================================================================================== #
if (i + 1) % self.n_critic == 0:
# Original-to-target domain.
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake)
g_loss_fake = -torch.mean(out_src)
g_loss_cls = self.classification_loss(out_cls, label_trg,
self.dataset)
# Target-to-original domain.
x_reconst = self.G(x_fake, c_org)
g_loss_rec = torch.mean(torch.abs(x_real - x_reconst))
# Backward and optimize.
g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + self.lambda_cls * g_loss_cls
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging.
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_cls'] = g_loss_cls.item()
# =================================================================================== #
# 4. Miscellaneous #
# =================================================================================== #
# Print out training information.
if (i + 1) % self.log_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
log = "Elapsed [{}], Iteration [{}/{}]".format(
et, i + 1, self.num_iters)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
# Translate fixed images for debugging.
if (i + 1) % self.sample_step == 0:
with torch.no_grad():
x_fake_list = [x_fixed]
for c_fixed in c_fixed_list:
x_fake_list.append(self.G(x_fixed, c_fixed))
x_concat = torch.cat(x_fake_list, dim=3)
sample_path = os.path.join(self.sample_dir,
'{}-images.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()),
sample_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(
sample_path))
# Save model checkpoints.
if (i + 1) % self.model_save_step == 0:
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(i + 1))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(i + 1))
torch.save(self.G.state_dict(), G_path)
torch.save(self.D.state_dict(), D_path)
print('Saved model checkpoints into {}...'.format(
self.model_save_dir))
# Decay learning rates.
if (i + 1) % self.lr_update_step == 0 and (i + 1) > (
self.num_iters - self.num_iters_decay):
g_lr -= (self.g_lr / float(self.num_iters_decay))
d_lr -= (self.d_lr / float(self.num_iters_decay))
self.update_lr(g_lr, d_lr)
print('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(
g_lr, d_lr))
def save_file(self, dir, image, type, num):
for k in range(len(image)):
image_k = [image[k]]
image_k = torch.cat(image_k, dim=2)
result_path = os.path.join(
dir, '{}-{}-images-{}.jpg'.format(num + 1, k + 1, type))
save_image(self.denorm(image_k.data.cpu()),
result_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(result_path))
def test(self):
"""Translate images using StarGAN trained on a single dataset."""
# Load the trained generator.
self.restore_model(self.test_iters)
# Set data loader.
if self.dataset == 'CelebA':
data_loader = self.celeba_loader
with torch.no_grad():
for i, (x_real, c_org) in enumerate(data_loader):
# Prepare input images and target domain labels.
x_real = x_real.to(self.device)
c_trg_list = self.create_labels(c_org, self.c_dim,
self.dataset,
self.selected_attrs)
# Translate images.
for j, c_trg in enumerate(c_trg_list):
# x_fake_list = []
# x_fake_list.append(self.G(x_real, c_trg))
x_fake = self.G(x_real, c_trg)
if j == 0:
self.save_file(self.result_dir, x_fake, 'black', j)
elif j == 1:
self.save_file(self.result_dir, x_fake, 'blond', j)
elif j == 2:
self.save_file(self.result_dir, x_fake, 'brown', j)
elif j == 3:
self.save_file(self.result_dir, x_fake, 'gender', j)
elif j == 4:
self.save_file(self.result_dir, x_fake, 'age', j)
def train(
dataset,
train_loader,
checkpoint_dir,
log_event_path,
nepochs,
learning_rate,
eval_per_step,
generator_step,
discriminator_step,
lambda_adv,
checkpoint_path,
seed,
):
torch.manual_seed(seed)
device = torch.device("cuda" if use_cuda else "cpu")
criterion = Loss(device, **loss_config)
# Model
model = Model(**network_config["nsf_config"]).to(device)
discriminator = Discriminator(
**network_config["discriminator_config"]).to(device)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
discriminator_optim = optim.Adam(discriminator.parameters(),
lr=learning_rate)
writer = SummaryWriter(log_event_path)
# train
epoch = 1
total_step = 0
current_lr = learning_rate
os.makedirs(checkpoint_dir, exist_ok=True)
if checkpoint_path != "":
model, discriminator, total_step, epoch = load_checkpoint(
checkpoint_path, model, optimizer, discriminator,
discriminator_optim)
current_lr = optimizer.param_groups[0]["lr"]
while epoch <= nepochs:
running_loss = 0
print("{}epoch:".format(epoch))
for step, (wav, mel, f0) in tqdm(enumerate(train_loader)):
model.train()
discriminator.train()
# configから操作できるようにはしたい
if total_step > 0 and current_lr > 1e-6 and total_step % 100000 == 0:
current_lr = current_lr / 2
for g_param_group, d_param_group in zip(
optimizer.param_groups,
discriminator_optim.param_groups):
g_param_group["lr"] = current_lr
d_param_group["lr"] = current_lr
optimizer.zero_grad()
discriminator_optim.zero_grad()
wav, mel, f0 = wav.to(device), mel.to(device), f0.to(device)
# Generator
if (total_step < generator_step
or total_step > generator_step + discriminator_step):
outputs = model(mel, f0)
stft_loss = criterion.stft_loss(outputs[:, :wav.size(-1)], wav)
if total_step < generator_step:
loss = stft_loss
adv_loss = None
else:
adv = discriminator(outputs.unsqueeze(1))
adv_loss = criterion.adversarial_loss(adv)
loss = stft_loss + lambda_adv * adv_loss
loss.backward()
optimizer.step()
else:
loss = None
stft_loss = None
adv_loss = None
# Discriminator
if total_step > generator_step:
with torch.no_grad():
outputs = model(mel, f0)
real = discriminator(wav.unsqueeze(1))
fake = discriminator(outputs.unsqueeze(1).detach())
real_loss, fake_loss = criterion.discriminator_loss(real, fake)
dis_loss = real_loss + fake_loss
dis_loss.backward()
discriminator_optim.step()
else:
dis_loss = None
if loss is not None:
writer.add_scalar("loss", float(loss.item()), total_step)
writer.add_scalar("stft_loss", float(stft_loss.item()),
total_step)
if adv_loss is not None:
writer.add_scalar("adv_loss", float(adv_loss.item()),
total_step)
if dis_loss is not None:
writer.add_scalar("dis_loss", float(dis_loss.item()),
total_step)
writer.add_scalar("real_loss", float(real_loss.item()),
total_step)
writer.add_scalar("fake_loss", float(fake_loss.item()),
total_step)
writer.add_scalar("learning_rate", current_lr, total_step)
total_step += 1
# running_loss += loss.item()
if total_step % eval_per_step == 0:
idx = np.random.randint(0, len(dataset.val_wav))
eval_model(
total_step,
writer,
device,
model,
dataset.get_all_length_data(idx),
checkpoint_dir,
data_config["mel_config"],
)
save_checkpoint(
model,
optimizer,
discriminator,
discriminator_optim,
total_step,
checkpoint_dir,
epoch,
)
# averaged_loss = running_loss / (len(train_loader))
# writer.add_scalar("loss (per epoch)", averaged_loss, epoch)
# print("Loss: {}".format(running_loss / (len(train_loader))))
epoch += 1
gamma = args.gamma
lambda_k = args.lambda_k
epochs = args.epochs
k = 0.0
transform = transforms.Compose([
transforms.Resize(img_size),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])])
train_data = datasets.MNIST( "../../data/mnist", train= True, download= True, transform= transform)
train_loader = torch.utils.data.DataLoader(train_data, batch_size= batch_size, shuffle= True)
#Load generator and discriminator
G = Generator(h, n).to(device)
D = Discriminator(h, n).to(device)
#Adam optimizer
optimizerG = optim.Adam(G.parameters(), lr=lr, betas=(beta1, beta2))
optimizerD = optim.Adam(D.parameters(), lr=lr, betas=(beta1, beta2))
#Training
G.train()
D.train()
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
for epoch in range(epochs):
for i, (real_images, _) in enumerate(train_loader):
real_images = Variable(real_images.type(Tensor))
#Train G
optimizerG.zero_grad()
left = chainer.as_variable(xp.array(lefteye).astype(xp.float32)).reshape(
1, 3, 32, 32)
right = chainer.as_variable(xp.array(righteye).astype(xp.float32)).reshape(
1, 3, 32, 32)
left = F.tile(left, (framesize, 1, 1, 1))
right = F.tile(right, (framesize, 1, 1, 1))
encoder = Encoder()
encoder.to_gpu()
enc_opt = set_optimizer(encoder)
refine = Refine()
refine.to_gpu()
ref_opt = set_optimizer(refine)
discriminator = Discriminator()
discriminator.to_gpu()
dis_opt = set_optimizer(discriminator)
for epoch in range(epochs):
sum_gen_loss = 0
sum_dis_loss = 0
for batch in range(0, iterations, framesize):
input_box = []
target_box = []
opt_box = []
rnd = np.random.randint(image_len)
dir_path = image_path + image_list[rnd]
ta = np.random.choice(["lefteye", "righteye"])
for index in range(framesize):
filename1 = dir_path + "/" + ta + "_" + str(0) + ".png"
state_dict = fill_statedict(state_dict, g_ema.vars, size)
g.load_state_dict(state_dict)
latent_avg = torch.from_numpy(g_ema.vars['dlatent_avg'].value().eval())
ckpt = {'g_ema': state_dict, 'latent_avg': latent_avg}
if args.gen:
g_train = Generator(size, 512, 8)
g_train_state = g_train.state_dict()
g_train_state = fill_statedict(g_train_state, generator.vars, size)
ckpt['g'] = g_train_state
if args.disc:
disc = Discriminator(size)
d_state = disc.state_dict()
d_state = discriminator_fill_statedict(d_state, discriminator.vars,
size)
ckpt['d'] = d_state
name = os.path.splitext(os.path.basename(args.path))[0]
torch.save(ckpt, name + '.pt')
batch_size = {256: 16, 512: 9, 1024: 4}
n_sample = batch_size.get(size, 25)
g = g.to(device)
z = np.random.RandomState(0).randn(n_sample, 512).astype('float32')
def main():
parser = argparse.ArgumentParser(description='Chainer: DCGAN MNIST')
parser.add_argument('--batchsize',
'-b',
type=int,
default=20,
help='Number of images in each mini-batch')
parser.add_argument('--epoch',
'-e',
type=int,
default=100,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu',
'-g',
type=int,
default=0,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--out',
'-o',
default='result',
help='Directory to output the result')
parser.add_argument('--resume',
'-r',
default='',
help='Resume the training from snapshot')
parser.add_argument('--n_hidden',
'-n',
type=int,
default=128,
help='Number of hidden units (z)')
parser.add_argument('--seed',
type=int,
default=0,
help='Random seed of z at visualization stage')
parser.add_argument('--snapshot_interval',
type=int,
default=1000,
help='Interval of snapshot')
parser.add_argument('--display_interval',
type=int,
default=100,
help='Interval of displaying log to console')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# n_hidden: {}'.format(args.n_hidden))
print('# epoch: {}'.format(args.epoch))
print('')
# Set up a neural network to train
gen = Generator(n_hidden=args.n_hidden)
dis = Discriminator()
if args.gpu >= 0:
# Make a specified GPU current
chainer.cuda.get_device_from_id(args.gpu).use()
gen.to_gpu() # Copy the model to the GPU
dis.to_gpu()
# Setup an optimizer
def make_optimizer(model, alpha=0.0002, beta1=0.5):
optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec')
return optimizer
opt_gen = make_optimizer(gen)
opt_dis = make_optimizer(dis)
# Load the MNIST dataset
train = load_image()
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
# Set up a trainer
updater = DCGANUpdater(models=(gen, dis),
iterator=train_iter,
optimizer={
'gen': opt_gen,
'dis': opt_dis
},
device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
epoch_interval = (1, 'epoch')
display_interval = (args.display_interval, 'iteration')
# trainer.extend(extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
# trainer.extend(extensions.snapshot_object(gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
# trainer.extend(extensions.snapshot_object(dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(
extensions.snapshot(filename='snapshot_epoch_{.updater.epoch}.npz'),
trigger=epoch_interval)
trainer.extend(extensions.snapshot_object(
gen, 'gen_epoch_{.updater.epoch}.npz'),
trigger=epoch_interval)
trainer.extend(extensions.snapshot_object(
dis, 'dis_epoch_{.updater.epoch}.npz'),
trigger=epoch_interval)
trainer.extend(extensions.LogReport(trigger=display_interval))
trainer.extend(extensions.PrintReport([
'epoch',
'iteration',
'gen/loss',
'dis/loss',
]),
trigger=display_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(out_generated_image(gen, dis, 10, 10, args.seed, args.out),
trigger=epoch_interval)
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
trainer.run()
# Variables
input_nc = 3
output_nc = 3
lr = 0.0002
g_lr = 0.0001
d_lr = 0.0001
batch_size = 32
size = 256
dataset = "CelebA"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = "cpu"
#initializing the generator and the discriminator
G = Generator(input_nc, output_nc).to(device)
D = Discriminator(input_nc).to(device)
# optimization function
g_optimizer = optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))
d_optimizer = optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))
#Loading data
image_dir = "../../data/CelebA/celeba"
attr_dir = "../../data/CelebA/list_attr_celeba.csv"
selected_attrs = ['Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Male', 'Young']
crop_size = 178
image_size = 128
batch_size = 32
data_loader = get_loader(image_dir, attr_dir, selected_attrs, crop_size,
image_size, batch_size, "CelebA", "train")
def main():
if os.path.exists(f'{args.job_dir}/checkpoint/model_best.pt'):
best_model = torch.load(f'{args.job_dir}/checkpoint/model_best.pt', map_location=torch.device(f"cuda:{args.gpus[0]}"))
model = prune_resnet(args, best_model['state_dict_s'])
if not os.path.exists(f'{args.job_dir}/pruned.pt'):
checkpoint = utils.checkpoint(args)
writer_train = SummaryWriter(args.job_dir + '/run/train')
writer_test = SummaryWriter(args.job_dir + '/run/test')
start_epoch = 0
best_prec1 = 0.0
best_prec5 = 0.0
# Data loading
print('=> Preparing data..')
loader = cifar10(args)
# Create model
print('=> Building model...')
model_t = resnet_56().to(args.gpus[0])
# Load teacher model
ckpt_t = torch.load(args.teacher_dir, map_location=torch.device(f"cuda:{args.gpus[0]}"))
state_dict_t = ckpt_t['state_dict']
model_t.load_state_dict(state_dict_t)
model_t = model_t.to(args.gpus[0])
for para in list(model_t.parameters())[:-2]:
para.requires_grad = False
model_s = resnet_56_sparse().to(args.gpus[0])
model_dict_s = model_s.state_dict()#模型的状态字典
model_dict_s.update(state_dict_t)
model_s.load_state_dict(model_dict_s)
if len(args.gpus) != 1:
model_s = nn.DataParallel(model_s, device_ids=args.gpus)
model_d = Discriminator().to(args.gpus[0])
models = [model_t, model_s, model_d]
optimizer_d = optim.SGD(model_d.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
param_s = [param for name, param in model_s.named_parameters() if 'mask' not in name]
param_m = [param for name, param in model_s.named_parameters() if 'mask' in name]
optimizer_s = optim.SGD(param_s, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
optimizer_m = FISTA(param_m, lr=args.lr, gamma=args.sparse_lambda)
#step_size和gamma代表每走step_size个epoch，学习率衰减gamma倍，阶梯形式
scheduler_d = StepLR(optimizer_d, step_size=args.lr_decay_step, gamma=0.1)
scheduler_s = StepLR(optimizer_s, step_size=args.lr_decay_step, gamma=0.1)
scheduler_m = StepLR(optimizer_m, step_size=args.lr_decay_step, gamma=0.1)
resume = args.resume
if resume:
print('=> Resuming from ckpt {}'.format(resume))
ckpt = torch.load(resume, map_location=torch.device(f"cuda:{args.gpus[0]}"))
best_prec1 = ckpt['best_prec1']
start_epoch = ckpt['epoch']
model_s.load_state_dict(ckpt['state_dict_s'])
model_d.load_state_dict(ckpt['state_dict_d'])
optimizer_d.load_state_dict(ckpt['optimizer_d'])
optimizer_s.load_state_dict(ckpt['optimizer_s'])
optimizer_m.load_state_dict(ckpt['optimizer_m'])
scheduler_d.load_state_dict(ckpt['scheduler_d'])
scheduler_s.load_state_dict(ckpt['scheduler_s'])
scheduler_m.load_state_dict(ckpt['scheduler_m'])
print('=> Continue from epoch {}...'.format(start_epoch))
optimizers = [optimizer_d, optimizer_s, optimizer_m]
schedulers = [scheduler_d, scheduler_s, scheduler_m]
if args.test_only:
test_prec1, test_prec5 = test(args, loader.loader_test, model_s)
print('=> Test Prec@1: {:.2f}'.format(test_prec1))
return
for epoch in range(start_epoch, args.num_epochs):
for s in schedulers:
s.step(epoch)
train(args, loader.loader_train, models, optimizers, epoch, writer_train)
test_prec1, test_prec5 = test(args, loader.loader_test, model_s)
is_best = best_prec1 < test_prec1
best_prec1 = max(test_prec1, best_prec1)
best_prec5 = max(test_prec5, best_prec5)
model_state_dict = model_s.module.state_dict() if len(args.gpus) > 1 else model_s.state_dict()
state = {
'state_dict_s': model_state_dict,
'state_dict_d': model_d.state_dict(),
'best_prec1': best_prec1,
'best_prec5': best_prec5,
'optimizer_d': optimizer_d.state_dict(),
'optimizer_s': optimizer_s.state_dict(),
'optimizer_m': optimizer_m.state_dict(),
'scheduler_d': scheduler_d.state_dict(),
'scheduler_s': scheduler_s.state_dict(),
'scheduler_m': scheduler_m.state_dict(),
'epoch': epoch + 1
}
checkpoint.save_model(state, epoch + 1, is_best, False)
print(f"=> Best @prec1: {best_prec1:.3f} @prec5: {best_prec5:.3f}")
best_model = torch.load(f'{args.job_dir}/checkpoint/model_best.pt', map_location=torch.device(f"cuda:{args.gpus[0]}"))
#pruning
model = prune_resnet(args, best_model['state_dict_s'])
else:
print('Have prunde!')
class Solver(object):
"""docstring for Solver."""
def __init__(self, data_loader, config):
self.config = config
self.data_loader = data_loader
# Model configurations.
self.lambda_cycle = config.lambda_cycle
self.lambda_cls = config.lambda_cls
self.lambda_identity = config.lambda_identity
# Training configurations.
self.data_dir = config.data_dir
self.test_dir = config.test_dir
self.batch_size = config.batch_size
self.num_iters = config.num_iters
self.num_iters_decay = config.num_iters_decay
self.g_lr = config.g_lr
self.d_lr = config.d_lr
self.c_lr = config.c_lr
self.n_critic = config.n_critic
self.beta1 = config.beta1
self.beta2 = config.beta2
self.resume_iters = config.resume_iters
# Test configurations.
self.test_iters = config.test_iters
self.trg_speaker = ast.literal_eval(config.trg_speaker)
self.src_speaker = config.src_speaker
# Miscellaneous.
self.use_tensorboard = config.use_tensorboard
self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
self.spk_enc = LabelBinarizer().fit(speakers)
# Directories.
self.log_dir = config.log_dir
self.sample_dir = config.sample_dir
self.model_save_dir = config.model_save_dir
self.result_dir = config.result_dir
# Step size.
self.log_step = config.log_step
self.sample_step = config.sample_step
self.model_save_step = config.model_save_step
self.lr_update_step = config.lr_update_step
# Build the model and tensorboard.
self.build_model()
if self.use_tensorboard:
self.build_tensorboard()
def build_model(self):
self.G = Generator()
self.D = Discriminator()
self.C = DomainClassifier()
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])
self.c_optimizer = torch.optim.Adam(self.C.parameters(), self.c_lr,[self.beta1, self.beta2])
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.print_network(self.C, 'C')
self.G.to(self.device)
self.D.to(self.device)
self.C.to(self.device)
def print_network(self, model, name):
"""Print out the network information."""
num_params = 0
for p in model.parameters():
num_params += p.numel()
print(model)
print(name)
print("The number of parameters: {}".format(num_params))
def build_tensorboard(self):
"""Build a tensorboard logger."""
from logger import Logger
self.logger = Logger(self.log_dir)
def update_lr(self, g_lr, d_lr, c_lr):
"""Decay learning rates of the generator and discriminator and classifier."""
for param_group in self.g_optimizer.param_groups:
param_group['lr'] = g_lr
for param_group in self.d_optimizer.param_groups:
param_group['lr'] = d_lr
for param_group in self.c_optimizer.param_groups:
param_group['lr'] = c_lr
def train(self):
# Learning rate cache for decaying.
g_lr = self.g_lr
d_lr = self.d_lr
c_lr = self.c_lr
start_iters = 0
if self.resume_iters:
pass
norm = Normalizer()
data_iter = iter(self.data_loader)
print('Start training......')
start_time = datetime.now()
for i in range(start_iters, self.num_iters):
# =================================================================================== #
# 1. Preprocess input data #
# =================================================================================== #
# Fetch real images and labels.
try:
x_real, speaker_idx_org, label_org = next(data_iter)
except:
data_iter = iter(self.data_loader)
x_real, speaker_idx_org, label_org = next(data_iter)
# Generate target domain labels randomly.
rand_idx = torch.randperm(label_org.size(0))
label_trg = label_org[rand_idx]
speaker_idx_trg = speaker_idx_org[rand_idx]
x_real = x_real.to(self.device) # Input images.
label_org = label_org.to(self.device) # Original domain one-hot labels.
label_trg = label_trg.to(self.device) # Target domain one-hot labels.
speaker_idx_org = speaker_idx_org.to(self.device) # Original domain labels
speaker_idx_trg = speaker_idx_trg.to(self.device) #Target domain labels
# =================================================================================== #
# 2. Train the discriminator #
# =================================================================================== #
# Compute loss with real audio frame.
## Modified
CELoss = nn.BCELoss()
m = nn.Sigmoid()
cls_real = self.C(x_real).squeeze(1)
cls_loss_real = CELoss(input=m(cls_real), target=speaker_idx_org)
self.reset_grad()
cls_loss_real.backward()
self.c_optimizer.step()
# Logging.
loss = {}
loss['C/C_loss'] = cls_loss_real.item()
out_r = self.D(x_real, label_org)
# Compute loss with fake audio frame.
x_fake = self.G(x_real, label_trg)
out_f = self.D(x_fake.detach(), label_trg)
d_loss_t = F.binary_cross_entropy_with_logits(input=out_f,target=torch.zeros_like(out_f, dtype=torch.float)) + \
F.binary_cross_entropy_with_logits(input=out_r, target=torch.ones_like(out_r, dtype=torch.float))
## Modified
out_cls = self.C(x_fake).squeeze(1)
d_loss_cls = CELoss(input=m(out_cls), target=speaker_idx_trg)
# Compute loss for gradient penalty.
alpha = torch.rand(x_real.size(0), 1, 1, 1).to(self.device)
x_hat = (alpha * x_real.data + (1 - alpha) * x_fake.data).requires_grad_(True)
out_src = self.D(x_hat, label_trg)
d_loss_gp = self.gradient_penalty(out_src, x_hat)
d_loss = d_loss_t + self.lambda_cls * d_loss_cls + 5*d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# loss['D/d_loss_t'] = d_loss_t.item()
# loss['D/loss_cls'] = d_loss_cls.item()
# loss['D/D_gp'] = d_loss_gp.item()
loss['D/D_loss'] = d_loss.item()
# =================================================================================== #
# 3. Train the generator #
# =================================================================================== #
if (i+1) % self.n_critic == 0:
# Original-to-target domain.
x_fake = self.G(x_real, label_trg)
g_out_src = self.D(x_fake, label_trg)
g_loss_fake = F.binary_cross_entropy_with_logits(input=g_out_src, target=torch.ones_like(g_out_src, dtype=torch.float))
## Modified
out_cls = self.C(x_real).squeeze(1)
g_loss_cls = CELoss(input=m(out_cls), target=speaker_idx_org)
# Target-to-original domain.
x_reconst = self.G(x_fake, label_org)
g_loss_rec = F.l1_loss(x_reconst, x_real )
# Original-to-Original domain(identity).
x_fake_iden = self.G(x_real, label_org)
id_loss = F.l1_loss(x_fake_iden, x_real )
# Backward and optimize.
g_loss = g_loss_fake + self.lambda_cycle * g_loss_rec +\
self.lambda_cls * g_loss_cls + self.lambda_identity * id_loss
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging.
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_cls'] = g_loss_cls.item()
loss['G/loss_id'] = id_loss.item()
loss['G/g_loss'] = g_loss.item()
# =================================================================================== #
# 4. Miscellaneous #
# =================================================================================== #
# Print out training information.
if (i+1) % self.log_step == 0:
et = datetime.now() - start_time
et = str(et)[:-7]
log = "Elapsed [{}], Iteration [{}/{}]".format(et, i+1, self.num_iters)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, i+1)
# Translate fixed images for debugging.
if (i+1) % self.sample_step == 0:
with torch.no_grad():
d, speaker = TestSet(self.test_dir).test_data()
target = random.choice([x for x in speakers if x != speaker])
label_t = self.spk_enc.transform([target])[0]
label_t = np.asarray([label_t])
for filename, content in d.items():
f0 = content['f0']
ap = content['ap']
sp_norm_pad = self.pad_coded_sp(content['coded_sp_norm'])
convert_result = []
for start_idx in range(0, sp_norm_pad.shape[1] - FRAMES + 1, FRAMES):
one_seg = sp_norm_pad[:, start_idx : start_idx+FRAMES]
one_seg = torch.FloatTensor(one_seg).to(self.device)
one_seg = one_seg.view(1,1,one_seg.size(0), one_seg.size(1))
l = torch.FloatTensor(label_t)
one_seg = one_seg.to(self.device)
l = l.to(self.device)
one_set_return = self.G(one_seg, l).data.cpu().numpy()
one_set_return = np.squeeze(one_set_return)
one_set_return = norm.backward_process(one_set_return, target)
convert_result.append(one_set_return)
convert_con = np.concatenate(convert_result, axis=1)
convert_con = convert_con[:, 0:content['coded_sp_norm'].shape[1]]
contigu = np.ascontiguousarray(convert_con.T, dtype=np.float64)
decoded_sp = decode_spectral_envelope(contigu, SAMPLE_RATE, fft_size=FFTSIZE)
f0_converted = norm.pitch_conversion(f0, speaker, target)
wav = synthesize(f0_converted, decoded_sp, ap, SAMPLE_RATE).astype(np.float32)
name = f'{speaker}-{target}_iter{i+1}_{filename}'
path = os.path.join(self.sample_dir, name)
print(f'[save]:{path}')
scipy.io.wavfile.write(path, SAMPLE_RATE, wav)
# Save model checkpoints.
if (i+1) % self.model_save_step == 0:
G_path = os.path.join(self.model_save_dir, '{}-G.ckpt'.format(i+1))
D_path = os.path.join(self.model_save_dir, '{}-D.ckpt'.format(i+1))
C_path = os.path.join(self.model_save_dir, '{}-C.ckpt'.format(i+1))
torch.save(self.G.state_dict(), G_path)
torch.save(self.D.state_dict(), D_path)
torch.save(self.C.state_dict(), C_path)
print('Saved model checkpoints into {}...'.format(self.model_save_dir))
# Decay learning rates.
if (i+1) % self.lr_update_step == 0 and (i+1) > (self.num_iters - self.num_iters_decay):
g_lr -= (self.g_lr / float(self.num_iters_decay))
d_lr -= (self.d_lr / float(self.num_iters_decay))
c_lr -= (self.c_lr / float(self.num_iters_decay))
self.update_lr(g_lr, d_lr, c_lr)
print ('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(g_lr, d_lr))
def gradient_penalty(self, y, x):
"""Compute gradient penalty: (L2_norm(dy/dx) - 1)**2."""
weight = torch.ones(y.size()).to(self.device)
dydx = torch.autograd.grad(outputs=y,
inputs=x,
grad_outputs=weight,
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
dydx = dydx.view(dydx.size(0), -1)
dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
return torch.mean((dydx_l2norm-1)**2)
def reset_grad(self):
"""Reset the gradient buffers."""
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
self.c_optimizer.zero_grad()
def restore_model(self, resume_iters):
"""Restore the trained generator and discriminator."""
print('Loading the trained models from step {}...'.format(resume_iters))
G_path = os.path.join(self.model_save_dir, '{}-G.ckpt'.format(resume_iters))
D_path = os.path.join(self.model_save_dir, '{}-D.ckpt'.format(resume_iters))
C_path = os.path.join(self.model_save_dir, '{}-C.ckpt'.format(resume_iters))
self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))
self.D.load_state_dict(torch.load(D_path, map_location=lambda storage, loc: storage))
self.C.load_state_dict(torch.load(C_path, map_location=lambda storage, loc: storage))
@staticmethod
def pad_coded_sp(coded_sp_norm):
f_len = coded_sp_norm.shape[1]
if f_len >= FRAMES:
pad_length = FRAMES-(f_len - (f_len//FRAMES) * FRAMES)
elif f_len < FRAMES:
pad_length = FRAMES - f_len
sp_norm_pad = np.hstack((coded_sp_norm, np.zeros((coded_sp_norm.shape[0], pad_length))))
return sp_norm_pad
def test(self):
"""Translate speech using StarGAN ."""
# Load the trained generator.
self.restore_model(self.test_iters)
norm = Normalizer()
# Set data loader.
d, speaker = TestSet(self.test_dir).test_data(self.src_speaker)
targets = self.trg_speaker
for target in targets:
print(target)
assert target in speakers
label_t = self.spk_enc.transform([target])[0]
label_t = np.asarray([label_t])
with torch.no_grad():
for filename, content in d.items():
f0 = content['f0']
ap = content['ap']
sp_norm_pad = self.pad_coded_sp(content['coded_sp_norm'])
convert_result = []
for start_idx in range(0, sp_norm_pad.shape[1] - FRAMES + 1, FRAMES):
one_seg = sp_norm_pad[:, start_idx : start_idx+FRAMES]
one_seg = torch.FloatTensor(one_seg).to(self.device)
one_seg = one_seg.view(1,1,one_seg.size(0), one_seg.size(1))
l = torch.FloatTensor(label_t)
one_seg = one_seg.to(self.device)
l = l.to(self.device)
one_set_return = self.G(one_seg, l).data.cpu().numpy()
one_set_return = np.squeeze(one_set_return)
one_set_return = norm.backward_process(one_set_return, target)
convert_result.append(one_set_return)
convert_con = np.concatenate(convert_result, axis=1)
convert_con = convert_con[:, 0:content['coded_sp_norm'].shape[1]]
contigu = np.ascontiguousarray(convert_con.T, dtype=np.float64)
decoded_sp = decode_spectral_envelope(contigu, SAMPLE_RATE, fft_size=FFTSIZE)
f0_converted = norm.pitch_conversion(f0, speaker, target)
wav = synthesize(f0_converted, decoded_sp, ap, SAMPLE_RATE).astype(np.float32)
name = f'{speaker}-{target}_iter{self.test_iters}_{filename}'
path = os.path.join(self.result_dir, name)
print(f'[save]:{path}')
scipy.io.wavfile.write(path, SAMPLE_RATE, wav)
size = args.size
batchsize = args.batchsize
outdir = "./output"
if not os.path.exists(outdir):
os.mkdir(outdir)
image_path = "./syani"
image_list = os.listdir(image_path)
list_len = len(image_list)
generator = Generator()
generator.to_gpu()
gen_opt = set_optimizer(generator)
discriminator = Discriminator()
discriminator.to_gpu()
dis_opt = set_optimizer(discriminator)
ztest = xp.random.uniform(-1, 1, (batchsize, 128)).astype(xp.float32)
ztest = chainer.as_variable(ztest)
for epoch in range(epochs):
sum_gen_loss = 0
sum_dis_loss = 0
for batch in range(0, iterations, framesize):
batch_box = []
for _ in range(batchsize):
frame_box = []
start_frame = np.random.randint(1, 1000 - framesize)
for index in range(framesize):
class DCGAN:
def __init__(self, CONFIG):
self.batch_size = CONFIG['batch_size']
self.latent_input = CONFIG['latent_input']
self.nb_image_to_gen = CONFIG['nb_image_to_gen']
self.image_size = CONFIG['image_size']
self.image_channels = CONFIG['image_channels']
self.save_path = CONFIG['save_path']
self.packing = CONFIG['packing']
self.real_label_smoothing = bool(CONFIG['real_label_smoothing'])
self.fake_label_smoothing = bool(CONFIG['fake_label_smoothing'])
self.nb_discriminator_step = CONFIG['nb_discriminator_step']
self.nb_generator_step = CONFIG['nb_generator_step']
# Device (cpu or gpu)
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Models
self.generator = Generator(CONFIG['latent_input'], CONFIG['model_complexity'], CONFIG['dropout_prob'],
CONFIG['weights_mean'], CONFIG['weights_std'],
CONFIG['image_channels']).to(self.device)
self.discriminator = Discriminator(CONFIG['model_complexity'], CONFIG['weights_mean'], CONFIG['weights_std'],
CONFIG['packing'], CONFIG['image_channels']).to(self.device)
print("------- GENERATOR ---------")
print(self.generator)
print("------- DISCRIMINATOR ---------")
print(self.discriminator)
# Optimizers
self.D_optimiser = optim.Adam(self.discriminator.parameters(), lr=CONFIG['learning_rate'],
betas=(CONFIG['beta1'], CONFIG['beta2']))
self.G_optimiser = optim.Adam(self.generator.parameters(), lr=CONFIG['learning_rate'],
betas=(CONFIG['beta1'], CONFIG['beta2']))
self.generator_losses = []
self.discriminator_losses = []
self.saved_latent_input = torch.randn(
(CONFIG['nb_image_to_gen'] * CONFIG['nb_image_to_gen'], CONFIG['latent_input'], 1, 1)).to(self.device)
# Create directory for the results if it doesn't already exists
import os
os.makedirs(self.save_path, exist_ok=True)
os.makedirs(self.save_path + "real/", exist_ok=True)
def load_dataset(self):
image_size = 32
batch_size = 128
root = "../datasets/MNIST_data"
trans = transforms.Compose([
transforms.Resize(image_size),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
# Load dataset
train_set = datasets.MNIST(root=root, train=True, transform=trans, download=True)
print('Number of images: ', len(train_set))
print('Sample image shape: ', train_set[0][0].shape, end='\n\n')
self.train_loader = torch.utils.data.DataLoader(
dataset=train_set,
batch_size=batch_size,
shuffle=True)
def train(self, nb_epoch=CONFIG['nb_epoch']):
print("Start training.")
for epoch in range(nb_epoch):
print("Epoch : " + str(epoch))
g_loss = []
d_loss = []
for batch_id, (x, target) in enumerate(self.train_loader):
real_batch_data = x.to(self.device)
current_batch_size = x.shape[0]
packed_real_data = pack(real_batch_data, self.packing)
packed_batch_size = packed_real_data.shape[0]
# labels
label_real = torch.full((packed_batch_size,), 1, device=self.device).squeeze()
label_fake = torch.full((packed_batch_size,), 0, device=self.device).squeeze()
# smoothed real labels between 0.7 and 1, and fake between 0 and 0.3
label_real_smooth = torch.rand((packed_batch_size,)).to(self.device).squeeze() * 0.3 + 0.7
label_fake_smooth = torch.rand((packed_batch_size,)).to(self.device).squeeze() * 0.3
temp_discriminator_loss = []
temp_generator_loss = []
### Train discriminator multiple times
for i in range(self.nb_discriminator_step):
loss_discriminator_total = self.train_discriminator(packed_real_data,
current_batch_size,
label_real_smooth if self.real_label_smoothing else label_real,
label_fake_smooth if self.fake_label_smoothing else label_fake)
temp_discriminator_loss.append(loss_discriminator_total.item())
# print("Discriminator step ", str(i), " with loss : ", loss_discriminator_total.item())
### Train generator multiple times
for i in range(self.nb_generator_step):
loss_generator_total = self.train_generator(current_batch_size, label_real)
temp_generator_loss.append(loss_generator_total.item())
if batch_id == len(self.train_loader) - 2:
save_images(real_batch_data, self.save_path + "real/", self.image_size, self.image_channels,
self.nb_image_to_gen, epoch)
### Keep track of losses
d_loss.append(torch.mean(torch.tensor(temp_discriminator_loss)))
g_loss.append(torch.mean(torch.tensor(temp_generator_loss)))
self.discriminator_losses.append(torch.mean(torch.tensor(d_loss)))
self.generator_losses.append(torch.mean(torch.tensor(g_loss)))
save_images(self.generator(self.saved_latent_input), self.save_path + "gen_", self.image_size,
self.image_channels, self.nb_image_to_gen, epoch)
write_loss_plot(self.generator_losses, "G loss", self.save_path, clear_plot=False)
write_loss_plot(self.discriminator_losses, "D loss", self.save_path, clear_plot=True)
print("Training finished.")
def train_discriminator(self, real_data, current_batch_size, real_label, fake_label):
# Generate with noise
latent_noise = torch.randn(current_batch_size, self.latent_input, 1, 1, device=self.device)
generated_batch = self.generator(latent_noise)
fake_data = pack(generated_batch, self.packing)
### Train discriminator
self.discriminator.zero_grad()
# Train on real data
real_prediction = self.discriminator(real_data).squeeze()
loss_discriminator_real = self.discriminator.loss(real_prediction, real_label)
# loss_discriminator_real.backward()
# Train on fake data
fake_prediction = self.discriminator(fake_data.detach()).squeeze()
loss_discriminator_fake = self.discriminator.loss(fake_prediction, fake_label)
# loss_discriminator_fake.backward()
# Add losses
loss_discriminator_total = loss_discriminator_real + loss_discriminator_fake
loss_discriminator_total.backward()
self.D_optimiser.step()
return loss_discriminator_total
def train_generator(self, current_batch_size, real_label):
# Generate with noise
latent_noise = torch.randn(current_batch_size, self.latent_input, 1, 1, device=self.device)
generated_batch = self.generator(latent_noise)
fake_data = pack(generated_batch, self.packing)
### Train generator
self.generator.zero_grad()
fake_prediction = self.discriminator(fake_data).squeeze()
# Loss
loss_generator = self.generator.loss(fake_prediction, real_label)
loss_generator.backward()
self.G_optimiser.step()
return loss_generator
def save_models(self):
save_model(self.generator, self.save_path, "generator_end")
save_model(self.discriminator, self.save_path, "discriminator_end")
def train(model, epochs, method="all_node", ablation="all"):
# 提前加载训练集和验证集到内存，节约时间。
def index_to_feature_wrapper(dict):
return mp.index_to_features(dict, data.x, method)
start_select = 50
train_index = train_list[:, 0].tolist()
print("Loading dataset with thread pool...")
train_metapath = pool.map(node_search_wrapper, train_index)
train_features = pool.map(index_to_feature_wrapper, train_metapath)
val_index = val_list[:, 0].tolist()
val_label = val_list[:, 1]
val_metapath = pool.map(node_search_wrapper, val_index)
val_features = pool.map(index_to_feature_wrapper, val_metapath)
lr = learning_rate
model.train() # 训练模式
best_micro_f1 = 0
best_macro_f1 = 0
type_set = set()
metapath_set = {}
for node in val_metapath:
for key in node:
type_set.add(key[0])
for type in type_set:
metapath_set[type] = set()
for node in val_metapath:
for key in node:
if len(node) - 1 > len(metapath_set[key[0]]):
if len(key) > 1:
metapath_set[key[0]].add(key)
metapath_label = {}
metapath_onehot = {}
discriminator = {}
d_optimizer = {}
label = {}
for type in type_set:
metapath_label[type] = {}
metapath_onehot[type] = {}
label[type] = []
for i, metapath in enumerate(metapath_set[type]):
metapath_label[type][metapath] = torch.zeros(batch_size,
data.type_num,
device=DEVICE)
metapath_onehot[type][metapath] = torch.zeros(
batch_size, device=DEVICE).long()
metapath_onehot[type][metapath][:] = i
for all_type in data.node_dict:
if all_type in metapath[1:]:
metapath_label[type][metapath][:, data.
node_dict[all_type]] = 1
label[type].append(metapath_label[type][metapath])
label[type] = torch.cat(label[type], dim=0)
discriminator[type] = Discriminator(info_section,
data.type_num).to(DEVICE)
d_optimizer[type] = optim.Adam(discriminator[type].parameters(),
lr=0.01,
weight_decay=5e-4)
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=5e-4)
select_flag = False
time1 = time.time()
for e in range(epochs):
# if single_path_limit is not None and (e + 1) % 20 == 0:
# print("Re-sampling...")
# train_metapath = pool.map(node_search_wrapper, train_index)
# train_features = pool.map(index_to_feature_wrapper, train_metapath)
for batch in range(num_batch_per_epoch):
batch_src_choice = np.random.choice(range(train_list.shape[0]),
size=(batch_size, ),
replace=False)
batch_src_index = train_list[batch_src_choice, 0]
batch_src_label = train_list[batch_src_choice, 1]
batch_feature_list = [train_features[i] for i in batch_src_choice]
batch_train_feature_dict, batch_src_index_dict, batch_src_label_dict, batch_train_rows_dict = mp.combine_features_dict(
batch_feature_list, batch_src_index, batch_src_label, DEVICE)
optimizer.zero_grad()
for type in d_optimizer:
d_optimizer[type].zero_grad()
if e >= start_select:
batch_train_logits_dict, GAN_input = model(
batch_train_rows_dict, batch_train_feature_dict, True)
else:
batch_train_logits_dict, GAN_input = model(
batch_train_rows_dict, batch_train_feature_dict,
False) # 获取模型的输出
for type in batch_train_logits_dict:
Loss_Classification = criterion(batch_train_logits_dict[type],
batch_src_label_dict[type])
assert ablation in ["all", "no_align"]
if ablation == "all":
Pred_D = discriminator[type](GAN_input[type])
Pred_Shuffle = discriminator[type](GAN_input[type], True)
Sorted_Pred_D = []
Sorted_Pred_Shuffle = []
for metapath in metapath_set[type]:
Sorted_Pred_D.append(Pred_D[metapath])
Sorted_Pred_Shuffle.append(Pred_Shuffle[metapath])
Sorted_Pred_D = torch.cat(Sorted_Pred_D, dim=0)
Sorted_Pred_Shuffle = torch.cat(Sorted_Pred_Shuffle, dim=0)
Loss_D = nn.BCELoss()(Sorted_Pred_D, label[type])
Loss_D_Shuffle = nn.BCELoss()(Sorted_Pred_Shuffle,
torch.zeros_like(
Sorted_Pred_Shuffle,
device=DEVICE))
Loss = Loss_Classification + Loss_D + Loss_D_Shuffle
else:
Loss = Loss_Classification
Loss.backward()
d_optimizer[type].step()
optimizer.step()
if e >= start_select and select_flag == False:
select_flag = True
pretrain_convergence = time2 - time1
print("Start select! Best f1-score reset to 0.")
print("Pretrain convergence time:", pretrain_convergence)
time1 = time.time()
best_micro_f1 = 0
best_macro_f1 = 0
if select_flag:
micro_f1, macro_f1 = val(model, val_features, val_index, val_label,
True)
model.show_metapath_importance()
else:
micro_f1, macro_f1 = val(model, val_features, val_index, val_label)
if micro_f1 >= best_micro_f1:
if micro_f1 > best_micro_f1:
time2 = time.time()
best_micro_f1 = micro_f1
best_macro_f1 = macro_f1
if select_flag:
torch.save(model.state_dict(),
"checkpoint/" + dataset + "_best_val")
elif macro_f1 > best_macro_f1:
best_micro_f1 = micro_f1
best_macro_f1 = macro_f1
if select_flag:
torch.save(model.state_dict(),
"checkpoint/" + dataset + "_best_val")
select_convergence = time2 - time1
print("Epoch ", e, ",Val Micro_f1 is ", micro_f1, ", Macro_f1 is ",
macro_f1, ", the best micro is ", best_micro_f1,
", the best macro is ", best_macro_f1)
torch.save(model.state_dict(), "checkpoint/" + dataset + "_final")
shuffle=False,
pin_memory=True)
validloader = DataLoader(valid_dataset,
# batch_size=int(batch_size / 2),
batch_size=batch_size,
num_workers=4,
shuffle=False,
pin_memory=True)
# Loss function
adversarial_loss = torch.nn.BCELoss()
# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()
if torch.cuda.is_available():
generator = generator.to('cuda')
discriminator = discriminator.to('cuda')
adversarial_loss = adversarial_loss.to('cuda')
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(), lr=1e-4, betas=(0.5, 0.999))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=1e-4, betas=(0.5, 0.999))
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
for epoch in range(epochs):
for i, (img, label) in enumerate(trainloader):
# print(img.size(1))
choices=['wgan-gp', 'r1'],
help='class of gan loss',
)
parser.add_argument(
'-d',
'--data',
default='folder',
type=str,
choices=['folder', 'lsun'],
help=('Specify dataset. ' 'Currently Image Folder and LSUN is supported'),
)
args = parser.parse_args()
generator = nn.DataParallel(StyledGenerator(code_size)).cuda()
discriminator = nn.DataParallel(Discriminator()).cuda()
g_running = StyledGenerator(code_size).cuda()
g_running.train(False)
class_loss = nn.CrossEntropyLoss()
g_optimizer = optim.Adam(
generator.module.generator.parameters(), lr=args.lr, betas=(0.0, 0.99)
)
g_optimizer.add_param_group(
{
'params': generator.module.style.parameters(),
'lr': args.lr * 0.01,
'mult': 0.01,
}
)
parser.add_argument('--mixing',
action='store_true',
help='use mixing regularization')
parser.add_argument(
'--loss',
type=str,
default='wgan-gp',
choices=['wgan-gp', 'r1'],
help='class of gan loss',
)
args = parser.parse_args()
generator = nn.DataParallel(StyledGenerator(code_size)).cuda()
discriminator = nn.DataParallel(
Discriminator(from_rgb_activate=not args.no_from_rgb_activate)).cuda()
g_running = StyledGenerator(code_size).cuda()
g_running.train(False)
g_optimizer = optim.Adam(generator.module.generator.parameters(),
lr=args.lr,
betas=(0.0, 0.99))
g_optimizer.add_param_group({
'params': generator.module.style.parameters(),
'lr': args.lr * 0.01,
'mult': 0.01,
})
d_optimizer = optim.Adam(discriminator.parameters(),
lr=args.lr,
betas=(0.0, 0.99))
upscale_factor=UPSCALE_FACTOR)
val_set = ValDatasetFromFolder('data/VOC2012/val',
upscale_factor=UPSCALE_FACTOR)
train_loader = DataLoader(dataset=train_set,
num_workers=4,
batch_size=64,
shuffle=True)
val_loader = DataLoader(dataset=val_set,
num_workers=4,
batch_size=1,
shuffle=False)
netG = Generator(UPSCALE_FACTOR)
print('# generator parameters:',
sum(param.numel() for param in netG.parameters()))
netD = Discriminator()
print('# discriminator parameters:',
sum(param.numel() for param in netD.parameters()))
generator_criterion = GeneratorLoss()
if torch.cuda.is_available():
netG.cuda()
netD.cuda()
generator_criterion.cuda()
optimizerG = optim.Adam(netG.parameters())
optimizerD = optim.Adam(netD.parameters())
results = {
'd_loss': [],
trans = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, ))])
MNIST_data = MNIST('./data', True, transform=trans, download=True)
loader = DataLoader(MNIST_data, BATCH_SIZE, True, num_workers=1)
# ~~~~~~~~~~~~~~~~~~~ creating tensorboard variables ~~~~~~~~~~~~~~~~~~~ #
writer_fake = SummaryWriter("logs/fake")
writer_real = SummaryWriter("logs/real")
# ~~~~~~~~~~~~~~~~~~~ loading the model ~~~~~~~~~~~~~~~~~~~ #
disc = Discriminator(in_features=CHANNELS, z_dim=Z_DIM).to(work_device)
gen = Faker(z_dim=Z_DIM).to(work_device)
# ~~~~~~~~~~~~~~~~~~~ create optimizer and loss ~~~~~~~~~~~~~~~~~~~ #
disc_optim = optim.Adam(disc.parameters(), lr)
gen_optim = optim.Adam(gen.parameters(), lr)
criterion = torch.nn.BCELoss()
# ~~~~~~~~~~~~~~~~~~~ training loop ~~~~~~~~~~~~~~~~~~~ #
for epoch in range(EPOCHS):
for batch_idx, (real, _) in enumerate(loader):
disc.train()
gen.train()
transforms.Normalize(
[0.5 for _ in range(CHANNELS_IMG)], [0.5 for _ in range(CHANNELS_IMG)]
),
]
)
# If you train on MNIST, remember to set channels_img to 1
dataset = datasets.MNIST(root="dataset/", train=True, transform=transforms,
download=True)
# comment mnist above and uncomment below if train on CelebA
#dataset = datasets.ImageFolder(root="celeb_dataset", transform=transforms)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
gen = Generator(NOISE_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)
disc = Discriminator(CHANNELS_IMG, FEATURES_DISC).to(device)
initialize_weights(gen)
initialize_weights(disc)
# print(gen) # 输出模型
# print(disc)
opt_gen = optim.Adam(gen.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))
opt_disc = optim.Adam(disc.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))
criterion = nn.BCELoss()
fixed_noise = torch.randn(32, NOISE_DIM, 1, 1).to(device)
writer_real = SummaryWriter(f"logs/real")
writer_fake = SummaryWriter(f"logs/fake")
step = 0
gen.train()
util.set_log_dir(args)
util.print_args(parser, args)
if args.arch == 'stylegan2':
from model import Generator, Discriminator
elif args.arch == 'swagan':
from swagan import Generator, Discriminator
generator = Generator(
args.size,
args.latent,
args.n_mlp_g,
channel_multiplier=args.channel_multiplier).to(device)
discriminator = Discriminator(args.size,
channel_multiplier=args.channel_multiplier,
which_phi=args.which_phi_d).to(device)
g_ema = Generator(args.size,
args.latent,
args.n_mlp_g,
channel_multiplier=args.channel_multiplier).to(device)
g_ema.eval()
accumulate(g_ema, generator, 0)
g_reg_ratio = args.g_reg_every / (args.g_reg_every +
1) if args.g_reg_every > 0 else 1.
d_reg_ratio = args.d_reg_every / (args.d_reg_every +
1) if args.d_reg_every > 0 else 1.
g_optim = optim.Adam(
generator.parameters(),
parser.add_argument('--dirA', required=True)
parser.add_argument('--dirB', required=True)
opt = parser.parse_args()
BETAS = (0.5, 0.999)
DECAY_EPOCH = opt.n_epochs // 2
dataset = ImagePairDataset(opt.dirA, opt.dirB)
dataloader = DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=4)
netG_A2B = Generator(3, 3).to(device)
netG_B2A = Generator(3, 3).to(device)
netD_A = Discriminator(3).to(device)
netD_B = Discriminator(3).to(device)
netG_A2B.apply(weights_init_normal)
netG_B2A.apply(weights_init_normal)
netD_A.apply(weights_init_normal)
netD_B.apply(weights_init_normal)
# optimizers and learning rate schedulers
optimizer_G = Adam(itertools.chain(netG_A2B.parameters(),
netG_B2A.parameters()),
lr=opt.lr,
betas=BETAS)
optimizer_D_en = Adam(netD_A.parameters(), lr=opt.lr, betas=BETAS)
optimizer_D_zh = Adam(netD_B.parameters(), lr=opt.lr, betas=BETAS)
class Solver(object):
"""Solver for training and testing StarGAN."""
def __init__(self, celeba_loader, config):
"""Initialize configurations."""
# Data loader.
self.celeba_loader = celeba_loader
# Model configurations.
self.c_dim = config.c_dim
self.image_size = config.image_size
self.g_conv_dim = config.g_conv_dim
self.d_conv_dim = config.d_conv_dim
self.g_repeat_num = config.g_repeat_num
self.d_repeat_num = config.d_repeat_num
self.lambda_cls = config.lambda_cls
self.lambda_rec = config.lambda_rec
# Training configurations.
self.dataset = config.dataset
self.batch_size = config.batch_size
self.num_iters = config.num_iters
self.num_iters_decay = config.num_iters_decay
self.g_lr = config.g_lr
self.d_lr = config.d_lr
self.beta1 = config.beta1
self.beta2 = config.beta2
self.resume_iters = config.resume_iters
self.selected_attrs = config.selected_attrs
# Test configurations.
self.test_iters = config.test_iters
# Miscellaneous.
self.use_tensorboard = config.use_tensorboard
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
# Directories.
self.log_dir = config.log_dir
self.sample_dir = config.sample_dir
self.model_save_dir = config.model_save_dir
self.result_dir = config.result_dir
# Step size.
self.log_step = config.log_step
self.fid_step = config.fid_step
self.sample_step = config.sample_step
self.model_save_step = config.model_save_step
self.lr_update_step = config.lr_update_step
# Classifier for computing FID
self.classifier = ResNet18().to(self.device)
self.classifier.eval()
# Build the model and tensorboard.
self.build_model()
if self.use_tensorboard:
self.build_tensorboard()
def build_model(self):
"""Create a generator and a discriminator."""
self.G = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.G_ema = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim,
self.d_repeat_num)
self.G_ema.load_state_dict(self.G.state_dict())
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr,
[self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr,
[self.beta1, self.beta2])
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.G.to(self.device)
self.G_ema.to(self.device)
self.D.to(self.device)
def compute_ema(self, beta=0.999):
for param, param_test in zip(self.G.parameters(),
self.G_ema.parameters()):
param_test.data = torch.lerp(param.data, param_test.data, beta)
for g_module, g_ema_module in zip(self.G.modules(),
self.G_ema.modules()):
if type(g_module) == nn.BatchNorm2d:
g_ema_module.running_mean = g_module.running_mean
g_ema_module.running_var = g_module.running_var
def print_network(self, model, name):
"""Print out the network information."""
num_params = 0
for p in model.parameters():
num_params += p.numel()
print(model)
print(name)
print("The number of parameters: {}".format(num_params))
def restore_model(self, resume_iters):
"""Restore the trained generator and discriminator."""
print(
'Loading the trained models from step {}...'.format(resume_iters))
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(resume_iters))
G_ema_path = os.path.join(self.model_save_dir,
'{}-G_ema.ckpt'.format(resume_iters))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(resume_iters))
self.G.load_state_dict(
torch.load(G_path, map_location=lambda storage, loc: storage))
self.G_ema.load_state_dict(
torch.load(G_ema_path, map_location=lambda storage, loc: storage))
self.D.load_state_dict(
torch.load(D_path, map_location=lambda storage, loc: storage))
def build_tensorboard(self):
"""Build a tensorboard logger."""
self.logger = Logger(self.log_dir)
def update_lr(self, g_lr, d_lr):
"""Decay learning rates of the generator and discriminator."""
for param_group in self.g_optimizer.param_groups:
param_group['lr'] = g_lr
for param_group in self.d_optimizer.param_groups:
param_group['lr'] = d_lr
def reset_grad(self):
"""Reset the gradient buffers."""
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
def denorm(self, x):
"""Convert the range from [-1, 1] to [0, 1]."""
out = (x + 1) / 2
return out.clamp_(0, 1)
def label2onehot(self, labels, dim):
"""Convert label indices to one-hot vectors."""
batch_size = labels.size(0)
out = torch.zeros(batch_size, dim)
out[np.arange(batch_size), labels.long()] = 1
return out
def create_labels(self, c_org, c_dim=5, selected_attrs=None):
"""Generate target domain labels for debugging and testing."""
# Get hair color indices.
hair_color_indices = []
for i, attr_name in enumerate(selected_attrs):
if attr_name in [
'Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Gray_Hair'
]:
hair_color_indices.append(i)
c_trg_list = []
for i in range(c_dim):
c_trg = c_org.clone()
if i in hair_color_indices: # Set one hair color to 1 and the rest to 0.
c_trg[:, i] = 1
for j in hair_color_indices:
if j != i:
c_trg[:, j] = 0
else:
c_trg[:, i] = (c_trg[:, i] == 0) # Reverse attribute value.
c_trg_list.append(c_trg.to(self.device))
return c_trg_list
def classification_loss(self, logit, target):
"""Compute binary or softmax cross entropy loss."""
return F.binary_cross_entropy_with_logits(
logit, target, size_average=False) / logit.size(0)
def train(self):
"""Train StarGAN within a single dataset."""
# Fetch fixed inputs for debugging.
data_iter = iter(self.celeba_loader)
x_fixed, c_org = next(data_iter)
x_fixed = x_fixed.to(self.device)
c_fixed_list = self.create_labels(c_org, self.c_dim,
self.selected_attrs)
# Learning rate cache for decaying.
g_lr = self.g_lr
d_lr = self.d_lr
# Start training from scratch or resume training.
start_iters = 0
if self.resume_iters:
start_iters = self.resume_iters
self.restore_model(self.resume_iters)
# Start training.
print('Start training...')
start_time = time.time()
for i in range(start_iters, self.num_iters):
# =================================================================================== #
# 1. Preprocess input data #
# =================================================================================== #
# Fetch real images and labels.
try:
x_real, label_org = next(data_iter)
except:
data_iter = iter(self.celeba_loader)
x_real, label_org = next(data_iter)
# Generate target domain labels randomly.
rand_idx = torch.randperm(label_org.size(0))
label_trg = label_org[rand_idx]
c_org = label_org.clone()
c_trg = label_trg.clone()
x_real = x_real.to(self.device) # Input images.
c_org = c_org.to(self.device) # Original domain labels.
c_trg = c_trg.to(self.device) # Target domain labels.
label_org = label_org.to(
self.device) # Labels for computing classification loss.
label_trg = label_trg.to(
self.device) # Labels for computing classification loss.
# =================================================================================== #
# 2. Train the discriminator #
# =================================================================================== #
# Compute loss with real images.
out_src, out_cls = self.D(x_real)
d_loss_real = F.mse_loss(
out_src, torch.ones_like(out_src, device=self.device))
d_loss_cls = self.classification_loss(out_cls, label_org)
# Compute loss with fake images.
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake.detach())
d_loss_fake = F.mse_loss(
out_src, torch.zeros_like(out_src, device=self.device))
# Backward and optimize.
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Logging.
loss = {
'D/loss_real': d_loss_real.item(),
'D/loss_fake': d_loss_fake.item(),
'D/loss_cls': d_loss_cls.item()
}
# =================================================================================== #
# 3. Train the generator #
# =================================================================================== #
# Original-to-target domain.
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake)
g_loss_fake = F.mse_loss(
out_src, torch.ones_like(out_src, device=self.device))
g_loss_cls = self.classification_loss(out_cls, label_trg)
# Target-to-original domain.
x_reconst = self.G(x_fake, c_org)
g_loss_rec = torch.mean(torch.abs(x_real - x_reconst))
# Backward and optimize.
g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + self.lambda_cls * g_loss_cls
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging.
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_cls'] = g_loss_cls.item()
# =================================================================================== #
# 4. Miscellaneous #
# =================================================================================== #
# Print out training information.
if (i + 1) % self.log_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
log = f'Elapsed [{et}], Iteration [{i + 1}/{self.num_iters}]'
for tag, value in loss.items():
log += f', {tag}: {value:.4f}'
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, i + 1)
# Translate fixed images for debugging.
if (i + 1) % self.sample_step == 0:
self.G_ema.eval()
with torch.no_grad():
x_fake_list = [x_fixed]
for c_fixed in c_fixed_list:
x_fake_list.append(self.G_ema(x_fixed, c_fixed))
x_concat = torch.cat(x_fake_list, dim=3)
sample_path = os.path.join(self.sample_dir,
f'{i + 1}-images.jpg')
save_image(self.denorm(x_concat.data.cpu()),
sample_path,
nrow=1,
padding=0)
print(f'Saved real and fake images into {sample_path}...')
# Save model checkpoints.
if (i + 1) % self.model_save_step == 0:
G_path = os.path.join(self.model_save_dir, f'{i + 1}-G.ckpt')
G_ema_path = os.path.join(self.model_save_dir,
f'{i + 1}-G_ema.ckpt')
D_path = os.path.join(self.model_save_dir, f'{i + 1}-D.ckpt')
torch.save(self.G.state_dict(), G_path)
torch.save(self.G_ema.state_dict(), G_ema_path)
torch.save(self.D.state_dict(), D_path)
print(f'Saved model checkpoints into {self.model_save_dir}...')
# Decay learning rates.
if (i + 1) % self.lr_update_step == 0 and (i + 1) > (
self.num_iters - self.num_iters_decay):
g_lr -= (self.g_lr / float(self.num_iters_decay))
d_lr -= (self.d_lr / float(self.num_iters_decay))
self.update_lr(g_lr, d_lr)
print(f'Decayed learning rates, g_lr: {g_lr}, d_lr: {d_lr}.')
# Count FID
if (i + 1) % self.fid_step == 0:
start = time.time()
fid = calculate_fid(self.celeba_loader, self.G_ema,
self.classifier)
print(f'FID score: {fid:.5f}, time: {time.time() - start:.5f}')
if self.use_tensorboard:
self.logger.scalar_summary('FID', fid, i + 1)
self.compute_ema()
@torch.no_grad()
def test(self):
"""Translate images using StarGAN trained on a single dataset."""
# Load the trained generator.
self.restore_model(self.test_iters)
data_loader = self.celeba_loader
for i, (x_real, c_org) in enumerate(data_loader):
# Prepare input images and target domain labels.
x_real = x_real.to(self.device)
c_trg_list = self.create_labels(c_org, self.c_dim,
self.selected_attrs)
# Translate images.
x_fake_list = [x_real]
for c_trg in c_trg_list:
x_fake_list.append(self.G_ema(x_real, c_trg))
# Save the translated images.
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.result_dir,
'{}-images.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()),
result_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(result_path))
from torch.utils.data import DataLoader
from loss import discriminatorLoss, generatorLoss
from dataset import TrainingDataset
from model import Generator, Discriminator
from trainer import Trainer
BATCH_SIZE = 64
dataset = TrainingDataset()
dataloader = DataLoader(dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=8)
model_G = Generator()
model_D = Discriminator()
MDL_OUTDIR = 'saved'
if not os.path.exists(MDL_OUTDIR):
os.mkdir(MDL_OUTDIR)
MDL_PRETRAINED_PATH_G = ''
MDL_PRETRAINED_PATH_D = ''
if MDL_PRETRAINED_PATH_G:
print('Loading pretrained model_G weights from:', MDL_PRETRAINED_PATH_G)
model_G.load_state_dict(torch.load(MDL_PRETRAINED_PATH_G))
else:
print('Initiating new model_G...')
if MDL_PRETRAINED_PATH_D:
print('Loading pretrained model_D weights from:', MDL_PRETRAINED_PATH_D)
model_D.load_state_dict(torch.load(MDL_PRETRAINED_PATH_D))
class Solver(object):
"""Solver for training and testing StarGAN."""
def __init__(self, celeba_loader, rafd_loader, config):
"""Initialize configurations."""
# Data loader.
self.celeba_loader = celeba_loader
self.rafd_loader = rafd_loader
# Model configurations.
self.c_dim = config.c_dim
self.c2_dim = config.c2_dim
self.image_size = config.image_size
self.g_conv_dim = config.g_conv_dim
self.d_conv_dim = config.d_conv_dim
self.g_repeat_num = config.g_repeat_num
self.d_repeat_num = config.d_repeat_num
self.lambda_cls = config.lambda_cls
self.lambda_rec = config.lambda_rec
self.lambda_gp = config.lambda_gp
# Training configurations.
self.dataset = config.dataset
self.batch_size = config.batch_size
self.num_iters = config.num_iters
self.num_iters_decay = config.num_iters_decay
self.g_lr = config.g_lr
self.d_lr = config.d_lr
self.n_critic = config.n_critic
self.beta1 = config.beta1
self.beta2 = config.beta2
self.resume_iters = config.resume_iters
self.selected_attrs = config.selected_attrs
# Test configurations.
self.test_iters = config.test_iters
# Miscellaneous.
self.use_tensorboard = config.use_tensorboard
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Directories.
self.log_dir = config.log_dir
self.sample_dir = config.sample_dir
self.model_save_dir = config.model_save_dir
self.result_dir = config.result_dir
# Step size.
self.log_step = config.log_step
self.sample_step = config.sample_step
self.model_save_step = config.model_save_step
self.lr_update_step = config.lr_update_step
# Build the model and tensorboard.
self.build_model()
if self.use_tensorboard:
self.build_tensorboard()
def build_model(self):
"""Create a generator and a discriminator."""
if self.dataset in ['CelebA', 'RaFD']:
self.G = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim, self.d_repeat_num)
elif self.dataset in ['Both']:
self.G = Generator(self.g_conv_dim, self.c_dim+self.c2_dim+2, self.g_repeat_num) # 2 for mask vector.
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim+self.c2_dim, self.d_repeat_num)
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.G.to(self.device)
self.D.to(self.device)
def print_network(self, model, name):
"""Print out the network information."""
num_params = 0
for p in model.parameters():
num_params += p.numel()
print(model)
print(name)
print("The number of parameters: {}".format(num_params))
def restore_model(self, resume_iters):
"""Restore the trained generator and discriminator."""
print('Loading the trained models from step {}...'.format(resume_iters))
G_path = os.path.join(self.model_save_dir, '{}-G.ckpt'.format(resume_iters))
D_path = os.path.join(self.model_save_dir, '{}-D.ckpt'.format(resume_iters))
self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))
self.D.load_state_dict(torch.load(D_path, map_location=lambda storage, loc: storage))
def build_tensorboard(self):
"""Build a tensorboard logger."""
from logger import Logger
self.logger = Logger(self.log_dir)
def update_lr(self, g_lr, d_lr):
"""Decay learning rates of the generator and discriminator."""
for param_group in self.g_optimizer.param_groups:
param_group['lr'] = g_lr
for param_group in self.d_optimizer.param_groups:
param_group['lr'] = d_lr
def reset_grad(self):
"""Reset the gradient buffers."""
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
def denorm(self, x):
"""Convert the range from [-1, 1] to [0, 1]."""
out = (x + 1) / 2
return out.clamp_(0, 1)
def gradient_penalty(self, y, x):
"""Compute gradient penalty: (L2_norm(dy/dx) - 1)**2."""
weight = torch.ones(y.size()).to(self.device)
dydx = torch.autograd.grad(outputs=y,
inputs=x,
grad_outputs=weight,
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
dydx = dydx.view(dydx.size(0), -1)
dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
return torch.mean((dydx_l2norm-1)**2)
def label2onehot(self, labels, dim):
"""Convert label indices to one-hot vectors."""
batch_size = labels.size(0)
out = torch.zeros(batch_size, dim)
out[np.arange(batch_size), labels.long()] = 1
return out
def create_labels(self, c_org, c_dim=5, dataset='CelebA', selected_attrs=None):
"""Generate target domain labels for debugging and testing."""
# Get hair color indices.
if dataset == 'CelebA':
hair_color_indices = []
for i, attr_name in enumerate(selected_attrs):
if attr_name in ['Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Gray_Hair']:
hair_color_indices.append(i)
c_trg_list = []
for i in range(c_dim):
if dataset == 'CelebA':
c_trg = c_org.clone()
if i in hair_color_indices: # Set one hair color to 1 and the rest to 0.
c_trg[:, i] = 1
for j in hair_color_indices:
if j != i:
c_trg[:, j] = 0
else:
c_trg[:, i] = (c_trg[:, i] == 0) # Reverse attribute value.
elif dataset == 'RaFD':
c_trg = self.label2onehot(torch.ones(c_org.size(0))*i, c_dim)
c_trg_list.append(c_trg.to(self.device))
return c_trg_list
def classification_loss(self, logit, target, dataset='CelebA'):
"""Compute binary or softmax cross entropy loss."""
if dataset == 'CelebA':
return F.binary_cross_entropy_with_logits(logit, target, size_average=False) / logit.size(0)
elif dataset == 'RaFD':
return F.cross_entropy(logit, target)
def train(self):
"""Train StarGAN within a single dataset."""
# Set data loader.
if self.dataset == 'CelebA':
data_loader = self.celeba_loader
elif self.dataset == 'RaFD':
data_loader = self.rafd_loader
# Fetch fixed inputs for debugging.
data_iter = iter(data_loader)
x_fixed, c_org = next(data_iter)
x_fixed = x_fixed.to(self.device)
c_fixed_list = self.create_labels(c_org, self.c_dim, self.dataset, self.selected_attrs)
# Learning rate cache for decaying.
g_lr = self.g_lr
d_lr = self.d_lr
# Start training from scratch or resume training.
start_iters = 0
if self.resume_iters:
start_iters = self.resume_iters
self.restore_model(self.resume_iters)
# Start training.
print('Start training...')
start_time = time.time()
for i in range(start_iters, self.num_iters):
# =================================================================================== #
# 1. Preprocess input data #
# =================================================================================== #
# Fetch real images and labels.
try:
x_real, label_org = next(data_iter)
except:
data_iter = iter(data_loader)
x_real, label_org = next(data_iter)
# Generate target domain labels randomly.
rand_idx = torch.randperm(label_org.size(0))
label_trg = label_org[rand_idx]
if self.dataset == 'CelebA':
c_org = label_org.clone()
c_trg = label_trg.clone()
elif self.dataset == 'RaFD':
c_org = self.label2onehot(label_org, self.c_dim)
c_trg = self.label2onehot(label_trg, self.c_dim)
x_real = x_real.to(self.device) # Input images.
c_org = c_org.to(self.device) # Original domain labels.
c_trg = c_trg.to(self.device) # Target domain labels.
label_org = label_org.to(self.device) # Labels for computing classification loss.
label_trg = label_trg.to(self.device) # Labels for computing classification loss.
# =================================================================================== #
# 2. Train the discriminator #
# =================================================================================== #
# Compute loss with real images.
out_src, out_cls = self.D(x_real)
d_loss_real = - torch.mean(out_src)
d_loss_cls = self.classification_loss(out_cls, label_org, self.dataset)
# Compute loss with fake images.
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake.detach())
d_loss_fake = torch.mean(out_src)
# Compute loss for gradient penalty.
alpha = torch.rand(x_real.size(0), 1, 1, 1).to(self.device)
x_hat = (alpha * x_real.data + (1 - alpha) * x_fake.data).requires_grad_(True)
out_src, _ = self.D(x_hat)
d_loss_gp = self.gradient_penalty(out_src, x_hat)
# Backward and optimize.
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls + self.lambda_gp * d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Logging.
loss = {}
loss['D/loss_real'] = d_loss_real.item()
loss['D/loss_fake'] = d_loss_fake.item()
loss['D/loss_cls'] = d_loss_cls.item()
loss['D/loss_gp'] = d_loss_gp.item()
# =================================================================================== #
# 3. Train the generator #
# =================================================================================== #
if (i+1) % self.n_critic == 0:
# Original-to-target domain.
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake)
g_loss_fake = - torch.mean(out_src)
g_loss_cls = self.classification_loss(out_cls, label_trg, self.dataset)
# Target-to-original domain.
x_reconst = self.G(x_fake, c_org)
g_loss_rec = torch.mean(torch.abs(x_real - x_reconst))
# Backward and optimize.
g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + self.lambda_cls * g_loss_cls
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging.
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_cls'] = g_loss_cls.item()
# =================================================================================== #
# 4. Miscellaneous #
# =================================================================================== #
# Print out training information.
if (i+1) % self.log_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
log = "Elapsed [{}], Iteration [{}/{}]".format(et, i+1, self.num_iters)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, i+1)
# Translate fixed images for debugging.
if (i+1) % self.sample_step == 0:
with torch.no_grad():
x_fake_list = [x_fixed]
for c_fixed in c_fixed_list:
x_fake_list.append(self.G(x_fixed, c_fixed))
x_concat = torch.cat(x_fake_list, dim=3)
sample_path = os.path.join(self.sample_dir, '{}-images.jpg'.format(i+1))
save_image(self.denorm(x_concat.data.cpu()), sample_path, nrow=1, padding=0)
print('Saved real and fake images into {}...'.format(sample_path))
# Save model checkpoints.
if (i+1) % self.model_save_step == 0:
G_path = os.path.join(self.model_save_dir, '{}-G.ckpt'.format(i+1))
D_path = os.path.join(self.model_save_dir, '{}-D.ckpt'.format(i+1))
torch.save(self.G.state_dict(), G_path)
torch.save(self.D.state_dict(), D_path)
print('Saved model checkpoints into {}...'.format(self.model_save_dir))
# Decay learning rates.
if (i+1) % self.lr_update_step == 0 and (i+1) > (self.num_iters - self.num_iters_decay):
g_lr -= (self.g_lr / float(self.num_iters_decay))
d_lr -= (self.d_lr / float(self.num_iters_decay))
self.update_lr(g_lr, d_lr)
print ('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(g_lr, d_lr))
def train_multi(self):
"""Train StarGAN with multiple datasets."""
# Data iterators.
celeba_iter = iter(self.celeba_loader)
rafd_iter = iter(self.rafd_loader)
# Fetch fixed inputs for debugging.
x_fixed, c_org = next(celeba_iter)
x_fixed = x_fixed.to(self.device)
c_celeba_list = self.create_labels(c_org, self.c_dim, 'CelebA', self.selected_attrs)
c_rafd_list = self.create_labels(c_org, self.c2_dim, 'RaFD')
zero_celeba = torch.zeros(x_fixed.size(0), self.c_dim).to(self.device) # Zero vector for CelebA.
zero_rafd = torch.zeros(x_fixed.size(0), self.c2_dim).to(self.device) # Zero vector for RaFD.
mask_celeba = self.label2onehot(torch.zeros(x_fixed.size(0)), 2).to(self.device) # Mask vector: [1, 0].
mask_rafd = self.label2onehot(torch.ones(x_fixed.size(0)), 2).to(self.device) # Mask vector: [0, 1].
# Learning rate cache for decaying.
g_lr = self.g_lr
d_lr = self.d_lr
# Start training from scratch or resume training.
start_iters = 0
if self.resume_iters:
start_iters = self.resume_iters
self.restore_model(self.resume_iters)
# Start training.
print('Start training...')
start_time = time.time()
for i in range(start_iters, self.num_iters):
for dataset in ['CelebA', 'RaFD']:
# =================================================================================== #
# 1. Preprocess input data #
# =================================================================================== #
# Fetch real images and labels.
data_iter = celeba_iter if dataset == 'CelebA' else rafd_iter
try:
x_real, label_org = next(data_iter)
except:
if dataset == 'CelebA':
celeba_iter = iter(self.celeba_loader)
x_real, label_org = next(celeba_iter)
elif dataset == 'RaFD':
rafd_iter = iter(self.rafd_loader)
x_real, label_org = next(rafd_iter)
# Generate target domain labels randomly.
rand_idx = torch.randperm(label_org.size(0))
label_trg = label_org[rand_idx]
if dataset == 'CelebA':
c_org = label_org.clone()
c_trg = label_trg.clone()
zero = torch.zeros(x_real.size(0), self.c2_dim)
mask = self.label2onehot(torch.zeros(x_real.size(0)), 2)
c_org = torch.cat([c_org, zero, mask], dim=1)
c_trg = torch.cat([c_trg, zero, mask], dim=1)
elif dataset == 'RaFD':
c_org = self.label2onehot(label_org, self.c2_dim)
c_trg = self.label2onehot(label_trg, self.c2_dim)
zero = torch.zeros(x_real.size(0), self.c_dim)
mask = self.label2onehot(torch.ones(x_real.size(0)), 2)
c_org = torch.cat([zero, c_org, mask], dim=1)
c_trg = torch.cat([zero, c_trg, mask], dim=1)
x_real = x_real.to(self.device) # Input images.
c_org = c_org.to(self.device) # Original domain labels.
c_trg = c_trg.to(self.device) # Target domain labels.
label_org = label_org.to(self.device) # Labels for computing classification loss.
label_trg = label_trg.to(self.device) # Labels for computing classification loss.
# =================================================================================== #
# 2. Train the discriminator #
# =================================================================================== #
# Compute loss with real images.
out_src, out_cls = self.D(x_real)
out_cls = out_cls[:, :self.c_dim] if dataset == 'CelebA' else out_cls[:, self.c_dim:]
d_loss_real = - torch.mean(out_src)
d_loss_cls = self.classification_loss(out_cls, label_org, dataset)
# Compute loss with fake images.
x_fake = self.G(x_real, c_trg)
out_src, _ = self.D(x_fake.detach())
d_loss_fake = torch.mean(out_src)
# Compute loss for gradient penalty.
alpha = torch.rand(x_real.size(0), 1, 1, 1).to(self.device)
x_hat = (alpha * x_real.data + (1 - alpha) * x_fake.data).requires_grad_(True)
out_src, _ = self.D(x_hat)
d_loss_gp = self.gradient_penalty(out_src, x_hat)
# Backward and optimize.
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls + self.lambda_gp * d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Logging.
loss = {}
loss['D/loss_real'] = d_loss_real.item()
loss['D/loss_fake'] = d_loss_fake.item()
loss['D/loss_cls'] = d_loss_cls.item()
loss['D/loss_gp'] = d_loss_gp.item()
# =================================================================================== #
# 3. Train the generator #
# =================================================================================== #
if (i+1) % self.n_critic == 0:
# Original-to-target domain.
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake)
out_cls = out_cls[:, :self.c_dim] if dataset == 'CelebA' else out_cls[:, self.c_dim:]
g_loss_fake = - torch.mean(out_src)
g_loss_cls = self.classification_loss(out_cls, label_trg, dataset)
# Target-to-original domain.
x_reconst = self.G(x_fake, c_org)
g_loss_rec = torch.mean(torch.abs(x_real - x_reconst))
# Backward and optimize.
g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + self.lambda_cls * g_loss_cls
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging.
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_cls'] = g_loss_cls.item()
# =================================================================================== #
# 4. Miscellaneous #
# =================================================================================== #
# Print out training info.
if (i+1) % self.log_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
log = "Elapsed [{}], Iteration [{}/{}], Dataset [{}]".format(et, i+1, self.num_iters, dataset)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, i+1)
# Translate fixed images for debugging.
if (i+1) % self.sample_step == 0:
with torch.no_grad():
x_fake_list = [x_fixed]
for c_fixed in c_celeba_list:
c_trg = torch.cat([c_fixed, zero_rafd, mask_celeba], dim=1)
x_fake_list.append(self.G(x_fixed, c_trg))
for c_fixed in c_rafd_list:
c_trg = torch.cat([zero_celeba, c_fixed, mask_rafd], dim=1)
x_fake_list.append(self.G(x_fixed, c_trg))
x_concat = torch.cat(x_fake_list, dim=3)
sample_path = os.path.join(self.sample_dir, '{}-images.jpg'.format(i+1))
save_image(self.denorm(x_concat.data.cpu()), sample_path, nrow=1, padding=0)
print('Saved real and fake images into {}...'.format(sample_path))
# Save model checkpoints.
if (i+1) % self.model_save_step == 0:
G_path = os.path.join(self.model_save_dir, '{}-G.ckpt'.format(i+1))
D_path = os.path.join(self.model_save_dir, '{}-D.ckpt'.format(i+1))
torch.save(self.G.state_dict(), G_path)
torch.save(self.D.state_dict(), D_path)
print('Saved model checkpoints into {}...'.format(self.model_save_dir))
# Decay learning rates.
if (i+1) % self.lr_update_step == 0 and (i+1) > (self.num_iters - self.num_iters_decay):
g_lr -= (self.g_lr / float(self.num_iters_decay))
d_lr -= (self.d_lr / float(self.num_iters_decay))
self.update_lr(g_lr, d_lr)
print ('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(g_lr, d_lr))
def test(self):
"""Translate images using StarGAN trained on a single dataset."""
# Load the trained generator.
self.restore_model(self.test_iters)
# Set data loader.
if self.dataset == 'CelebA':
data_loader = self.celeba_loader
elif self.dataset == 'RaFD':
data_loader = self.rafd_loader
with torch.no_grad():
for i, (x_real, c_org) in enumerate(data_loader):
# Prepare input images and target domain labels.
x_real = x_real.to(self.device)
c_trg_list = self.create_labels(c_org, self.c_dim, self.dataset, self.selected_attrs)
# Translate images.
x_fake_list = [x_real]
for c_trg in c_trg_list:
x_fake_list.append(self.G(x_real, c_trg))
# Save the translated images.
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.result_dir, '{}-images.jpg'.format(i+1))
save_image(self.denorm(x_concat.data.cpu()), result_path, nrow=1, padding=0)
print('Saved real and fake images into {}...'.format(result_path))
def test_multi(self):
"""Translate images using StarGAN trained on multiple datasets."""
# Load the trained generator.
self.restore_model(self.test_iters)
with torch.no_grad():
for i, (x_real, c_org) in enumerate(self.celeba_loader):
# Prepare input images and target domain labels.
x_real = x_real.to(self.device)
c_celeba_list = self.create_labels(c_org, self.c_dim, 'CelebA', self.selected_attrs)
c_rafd_list = self.create_labels(c_org, self.c2_dim, 'RaFD')
zero_celeba = torch.zeros(x_real.size(0), self.c_dim).to(self.device) # Zero vector for CelebA.
zero_rafd = torch.zeros(x_real.size(0), self.c2_dim).to(self.device) # Zero vector for RaFD.
mask_celeba = self.label2onehot(torch.zeros(x_real.size(0)), 2).to(self.device) # Mask vector: [1, 0].
mask_rafd = self.label2onehot(torch.ones(x_real.size(0)), 2).to(self.device) # Mask vector: [0, 1].
# Translate images.
x_fake_list = [x_real]
for c_celeba in c_celeba_list:
c_trg = torch.cat([c_celeba, zero_rafd, mask_celeba], dim=1)
x_fake_list.append(self.G(x_real, c_trg))
for c_rafd in c_rafd_list:
c_trg = torch.cat([zero_celeba, c_rafd, mask_rafd], dim=1)
x_fake_list.append(self.G(x_real, c_trg))
# Save the translated images.
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.result_dir, '{}-images.jpg'.format(i+1))
save_image(self.denorm(x_concat.data.cpu()), result_path, nrow=1, padding=0)
print('Saved real and fake images into {}...'.format(result_path))
class Solver(object):
def __init__(self, celeba_loader, rafd_loader, config):
self.celeba_loader = celeba_loader
self.rafd_loader = rafd_loader
#
self.c_dim = config.c_dim
self.c2_dim = config.c2_dim
self.image_size = config.image_size
self.g_conv_dim = config.g_conv_dim
self.d_conv_dim = config.d_conv_dim
self.g_repeat_num = config.g_repeat_num
self.d_repeat_num = config.d_repeat_num
self.lambda_cls = config.lambda_cls
self.lambda_rec = config.lambda_rec
self.lambda_gp = config.lambda_gp
self.dataset = config.dataset
self.batch_size = config.batch_size
self.num_iters = config.num_iters
self.num_iters_decay = config.num_iters_decay
self.g_lr = config.g_lr
self.d_lr = config.d_lr
self.n_critic = config.n_critic
self.beta1 = config.beta1
self.beta2 = config.beta2
self.resume_iters = config.resume_iters
self.selected_attrs = config.selected_attrs
self.test_iters = config.test_iters
self.use_tensorboard = config.use_tensorboard
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.log_dir = config.log_dir
self.sample_dir = config.sample_dir
self.model_save_dir = config.model_save_dir
self.result_dir = config.result_dir
self.log_step = config.log_step
self.sample_step = config.sample_step
self.model_save_step = config.model_save_step
self.lr_update_step = config.lr_update_step
self.build_model()
if self.use_tensorboard:
self.build_tensorboard()
def build_model(self):
if self.dataset in ['CelebA', 'RaFD']:
self.G = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.D = Discriminator(self.image_size, self.d_conv_dim,
self.c_dim, self.d_repeat_num)
elif self.dataset in ['Both']:
self.G = Generator(self.g_conv_dim, self.c_dim + self.c2_dim + 2,
self.g_repeat_num) # 2 for mask vector.
self.D = Discriminator(self.image_size, self.d_conv_dim,
self.c_dim + self.c2_dim, self.d_repeat_num)
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr,
[self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr,
[self.beta1, self.beta2])
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.G.to(self.device)
self.D.to(self.device)
def print_network(self, model, name):
num_params = 0
for p in model.parameters():
num_params += p.numel()
print(model)
print(name)
print("The number of parameters: {}".format(num_params))
def restore_model(self, resume_iters):
print(
'Loading the trained models from step {}...'.format(resume_iters))
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(resume_iters))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(resume_iters))
self.G.load_state_dict(
torch.load(G_path, map_location=lambda storage, loc: storage))
self.D.load_state_dict(
torch.load(D_path, map_location=lambda storage, loc: storage))
def build_tensorboard(self):
from logger import Logger
self.logger = Logger(self.log_dir)
def update_lr(self, g_lr, d_lr):
for param_group in self.g_optimizer.param_groups:
param_group['lr'] = g_lr
for param_group in self.d_optimizer.param_groups:
param_group['lr'] = d_lr
def reset_grad(self):
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
def denorm(self, x):
out = (x + 1) / 2
return out.clamp_(0, 1)
def gradient_penalty(self, y, x):
weight = torch.ones(y.size()).to(self.device)
dydx = torch.autograd.grad(outputs=y,
inputs=x,
grad_outputs=weight,
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
dydx = dydx.view(dydx.size(0), -1)
dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
return torch.mean((dydx_l2norm - 1)**2)
def label2onehot(self, labels, dim):
batch_size = labels.size(0)
out = torch.zeros(batch_size, dim)
out[np.arange(batch_size), labels.long()] = 1
return out
def create_labels(self,
c_org,
c_dim=5,
dataset='CelebA',
selected_attrs=None):
if dataset == 'CelebA':
hair_color_indices = []
for i, attr_name in enumerate(selected_attrs):
if attr_name in [
'Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Gray_Hair'
]:
hair_color_indices.append(i)
c_trg_list = []
for i in range(c_dim):
if dataset == 'CelebA':
c_trg = c_org.clone()
if i in hair_color_indices:
c_trg[:, i] = 1
for j in hair_color_indices:
if j != i:
c_trg[:, j] = 0
else:
c_trg[:, i] = (c_trg[:, i] == 0)
elif dataset == 'RaFD':
c_trg = self.label2onehot(torch.ones(c_org.size(0)) * i, c_dim)
c_trg_list.append(c_trg.to(self.device))
return c_trg_list
def classification_loss(self, logit, target, dataset='CelebA'):
if dataset == 'CelebA':
return F.binary_cross_entropy_with_logits(
logit, target, size_average=False) / logit.size(0)
elif dataset == 'RaFD':
return F.cross_entropy(logit, target)
def train(self):
# Set data loader.
if self.dataset == 'CelebA':
data_loader = self.celeba_loader
elif self.dataset == 'RaFD':
data_loader = self.rafd_loader
data_iter = iter(data_loader)
x_fixed, c_org = next(data_iter)
x_fixed = x_fixed.to(self.device)
c_fixed_list = self.create_labels(c_org, self.c_dim, self.dataset,
self.selected_attrs)
g_lr = self.g_lr
d_lr = self.d_lr
start_iters = 0
if self.resume_iters:
start_iters = self.resume_iters
self.restore_model(self.resume_iters)
# Start training.
print('Start training...')
start_time = time.time()
for i in range(start_iters, self.num_iters):
try:
x_real, label_org = next(data_iter)
except:
data_iter = iter(data_loader)
x_real, label_org = next(data_iter)
rand_idx = torch.randperm(label_org.size(0))
label_trg = label_org[rand_idx]
if self.dataset == 'CelebA':
c_org = label_org.clone()
c_trg = label_trg.clone()
elif self.dataset == 'RaFD':
c_org = self.label2onehot(label_org, self.c_dim)
c_trg = self.label2onehot(label_trg, self.c_dim)
x_real = x_real.to(self.device)
c_org = c_org.to(self.device)
c_trg = c_trg.to(self.device)
label_org = label_org.to(self.device)
label_trg = label_trg.to(self.device)
out_src, out_cls = self.D(x_real)
d_loss_real = -torch.mean(out_src)
d_loss_cls = self.classification_loss(out_cls, label_org,
self.dataset)
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake.detach())
d_loss_fake = torch.mean(out_src)
alpha = torch.rand(x_real.size(0), 1, 1, 1).to(self.device)
x_hat = (alpha * x_real.data +
(1 - alpha) * x_fake.data).requires_grad_(True)
out_src, _ = self.D(x_hat)
d_loss_gp = self.gradient_penalty(out_src, x_hat)
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls + self.lambda_gp * d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
loss = {}
loss['D/loss_real'] = d_loss_real.item()
loss['D/loss_fake'] = d_loss_fake.item()
loss['D/loss_cls'] = d_loss_cls.item()
loss['D/loss_gp'] = d_loss_gp.item()
if (i + 1) % self.n_critic == 0:
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake)
g_loss_fake = -torch.mean(out_src)
g_loss_cls = self.classification_loss(out_cls, label_trg,
self.dataset)
x_reconst = self.G(x_fake, c_org)
g_loss_rec = torch.mean(torch.abs(x_real - x_reconst))
g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + self.lambda_cls * g_loss_cls
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_cls'] = g_loss_cls.item()
if (i + 1) % self.log_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
log = "Elapsed [{}], Iteration [{}/{}]".format(
et, i + 1, self.num_iters)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, i + 1)
if (i + 1) % self.sample_step == 0:
with torch.no_grad():
x_fake_list = [x_fixed]
for c_fixed in c_fixed_list:
x_fake_list.append(self.G(x_fixed, c_fixed))
x_concat = torch.cat(x_fake_list, dim=3)
sample_path = os.path.join(self.sample_dir,
'{}-images.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()),
sample_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(
sample_path))
if (i + 1) % self.model_save_step == 0:
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(i + 1))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(i + 1))
torch.save(self.G.state_dict(), G_path)
torch.save(self.D.state_dict(), D_path)
print('Saved model checkpoints into {}...'.format(
self.model_save_dir))
if (i + 1) % self.lr_update_step == 0 and (i + 1) > (
self.num_iters - self.num_iters_decay):
g_lr -= (self.g_lr / float(self.num_iters_decay))
d_lr -= (self.d_lr / float(self.num_iters_decay))
self.update_lr(g_lr, d_lr)
print('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(
g_lr, d_lr))
def train_multi(self):
celeba_iter = iter(self.celeba_loader)
rafd_iter = iter(self.rafd_loader)
x_fixed, c_org = next(celeba_iter)
x_fixed = x_fixed.to(self.device)
c_celeba_list = self.create_labels(c_org, self.c_dim, 'CelebA',
self.selected_attrs)
c_rafd_list = self.create_labels(c_org, self.c2_dim, 'RaFD')
zero_celeba = torch.zeros(x_fixed.size(0), self.c_dim).to(
self.device) # Zero vector for CelebA.
zero_rafd = torch.zeros(x_fixed.size(0), self.c2_dim).to(
self.device) # Zero vector for RaFD.
mask_celeba = self.label2onehot(torch.zeros(x_fixed.size(0)), 2).to(
self.device) # Mask vector: [1, 0].
mask_rafd = self.label2onehot(torch.ones(x_fixed.size(0)), 2).to(
self.device) # Mask vector: [0, 1].
g_lr = self.g_lr
d_lr = self.d_lr
start_iters = 0
if self.resume_iters:
start_iters = self.resume_iters
self.restore_model(self.resume_iters)
print('Start training...')
start_time = time.time()
for i in range(start_iters, self.num_iters):
for dataset in ['CelebA', 'RaFD']:
data_iter = celeba_iter if dataset == 'CelebA' else rafd_iter
try:
x_real, label_org = next(data_iter)
except:
if dataset == 'CelebA':
celeba_iter = iter(self.celeba_loader)
x_real, label_org = next(celeba_iter)
elif dataset == 'RaFD':
rafd_iter = iter(self.rafd_loader)
x_real, label_org = next(rafd_iter)
rand_idx = torch.randperm(label_org.size(0))
label_trg = label_org[rand_idx]
if dataset == 'CelebA':
c_org = label_org.clone()
c_trg = label_trg.clone()
zero = torch.zeros(x_real.size(0), self.c2_dim)
mask = self.label2onehot(torch.zeros(x_real.size(0)), 2)
c_org = torch.cat([c_org, zero, mask], dim=1)
c_trg = torch.cat([c_trg, zero, mask], dim=1)
elif dataset == 'RaFD':
c_org = self.label2onehot(label_org, self.c2_dim)
c_trg = self.label2onehot(label_trg, self.c2_dim)
zero = torch.zeros(x_real.size(0), self.c_dim)
mask = self.label2onehot(torch.ones(x_real.size(0)), 2)
c_org = torch.cat([zero, c_org, mask], dim=1)
c_trg = torch.cat([zero, c_trg, mask], dim=1)
x_real = x_real.to(self.device)
c_org = c_org.to(self.device)
c_trg = c_trg.to(self.device)
label_org = label_org.to(self.device)
label_trg = label_trg.to(self.device)
out_src, out_cls = self.D(x_real)
out_cls = out_cls[:, :self.
c_dim] if dataset == 'CelebA' else out_cls[:,
self
.
c_dim:]
d_loss_real = -torch.mean(out_src)
d_loss_cls = self.classification_loss(out_cls, label_org,
dataset)
x_fake = self.G(x_real, c_trg)
out_src, _ = self.D(x_fake.detach())
d_loss_fake = torch.mean(out_src)
alpha = torch.rand(x_real.size(0), 1, 1, 1).to(self.device)
x_hat = (alpha * x_real.data +
(1 - alpha) * x_fake.data).requires_grad_(True)
out_src, _ = self.D(x_hat)
d_loss_gp = self.gradient_penalty(out_src, x_hat)
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls + self.lambda_gp * d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
loss = {}
loss['D/loss_real'] = d_loss_real.item()
loss['D/loss_fake'] = d_loss_fake.item()
loss['D/loss_cls'] = d_loss_cls.item()
loss['D/loss_gp'] = d_loss_gp.item()
if (i + 1) % self.n_critic == 0:
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake)
out_cls = out_cls[:, :self.
c_dim] if dataset == 'CelebA' else out_cls[:,
self
.
c_dim:]
g_loss_fake = -torch.mean(out_src)
g_loss_cls = self.classification_loss(
out_cls, label_trg, dataset)
x_reconst = self.G(x_fake, c_org)
g_loss_rec = torch.mean(torch.abs(x_real - x_reconst))
g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + self.lambda_cls * g_loss_cls
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_cls'] = g_loss_cls.item()
if (i + 1) % self.log_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
log = "Elapsed [{}], Iteration [{}/{}], Dataset [{}]".format(
et, i + 1, self.num_iters, dataset)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, i + 1)
if (i + 1) % self.sample_step == 0:
with torch.no_grad():
x_fake_list = [x_fixed]
for c_fixed in c_celeba_list:
c_trg = torch.cat([c_fixed, zero_rafd, mask_celeba],
dim=1)
x_fake_list.append(self.G(x_fixed, c_trg))
for c_fixed in c_rafd_list:
c_trg = torch.cat([zero_celeba, c_fixed, mask_rafd],
dim=1)
x_fake_list.append(self.G(x_fixed, c_trg))
x_concat = torch.cat(x_fake_list, dim=3)
sample_path = os.path.join(self.sample_dir,
'{}-images.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()),
sample_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(
sample_path))
if (i + 1) % self.model_save_step == 0:
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(i + 1))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(i + 1))
torch.save(self.G.state_dict(), G_path)
torch.save(self.D.state_dict(), D_path)
print('Saved model checkpoints into {}...'.format(
self.model_save_dir))
if (i + 1) % self.lr_update_step == 0 and (i + 1) > (
self.num_iters - self.num_iters_decay):
g_lr -= (self.g_lr / float(self.num_iters_decay))
d_lr -= (self.d_lr / float(self.num_iters_decay))
self.update_lr(g_lr, d_lr)
print('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(
g_lr, d_lr))
def test(self):
self.restore_model(self.test_iters)
if self.dataset == 'CelebA':
data_loader = self.celeba_loader
elif self.dataset == 'RaFD':
data_loader = self.rafd_loader
with torch.no_grad():
for i, (x_real, c_org) in enumerate(data_loader):
x_real = x_real.to(self.device)
c_trg_list = self.create_labels(c_org, self.c_dim,
self.dataset,
self.selected_attrs)
x_fake_list = [x_real]
for c_trg in c_trg_list:
x_fake_list.append(self.G(x_real, c_trg))
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.result_dir,
'{}-images.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()),
result_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(
result_path))
def test_multi(self):
self.restore_model(self.test_iters)
with torch.no_grad():
for i, (x_real, c_org) in enumerate(self.celeba_loader):
x_real = x_real.to(self.device)
c_celeba_list = self.create_labels(c_org, self.c_dim, 'CelebA',
self.selected_attrs)
c_rafd_list = self.create_labels(c_org, self.c2_dim, 'RaFD')
zero_celeba = torch.zeros(x_real.size(0), self.c_dim).to(
self.device) # Zero vector for CelebA.
zero_rafd = torch.zeros(x_real.size(0), self.c2_dim).to(
self.device) # Zero vector for RaFD.
mask_celeba = self.label2onehot(torch.zeros(
x_real.size(0)), 2).to(self.device) # Mask vector: [1, 0].
mask_rafd = self.label2onehot(torch.ones(
x_real.size(0)), 2).to(self.device) # Mask vector: [0, 1].
x_fake_list = [x_real]
for c_celeba in c_celeba_list:
c_trg = torch.cat([c_celeba, zero_rafd, mask_celeba],
dim=1)
x_fake_list.append(self.G(x_real, c_trg))
for c_rafd in c_rafd_list:
c_trg = torch.cat([zero_celeba, c_rafd, mask_rafd], dim=1)
x_fake_list.append(self.G(x_real, c_trg))
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.result_dir,
'{}-images.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()),
result_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(
result_path))
try:
os.mkdir(args.model_dir)
except:
pass
try:
os.mkdir(args.visualization_dir)
except:
pass
x_train, x_test, y_train, y_test = utils.load_mnist()
N = len(x_train)
model = AAE(784, n_z, hidden_units_enc=(1000, 1000, 500), hidden_units_dec=(500,1000,1000))
dis = Discriminator(n_z+10)
use_gpu = args.gpu >= 0
if use_gpu:
cuda.get_device(args.gpu).use()
model.to_gpu()
dis.to_gpu()
xp = np if args.gpu < 0 else cuda.cupy
optimizer_dis = optimizers.Adam(alpha=0.0002, beta1=0.5)
optimizer_aae = optimizers.Adam(alpha=0.0002, beta1=0.5)
optimizer_dis.setup(dis)
optimizer_aae.setup(model)
optimizer_dis.add_hook(optimizer.WeightDecay(0.0001))
optimizer_aae.add_hook(optimizer.WeightDecay(0.0001))
args.start_iter = 0
if args.arch == 'stylegan2':
from model import Generator, Discriminator
elif args.arch == 'swagan':
from swagan import Generator, Discriminator
generator = Generator(
args.size,
args.latent,
args.n_mlp,
channel_multiplier=args.channel_multiplier).to(device)
discriminator = Discriminator(args.size,
channel_multiplier=args.channel_multiplier,
patch_number=args.patch_number).to(device)
g_ema = Generator(args.size,
args.latent,
args.n_mlp,
channel_multiplier=args.channel_multiplier).to(device)
g_ema.eval()
accumulate(g_ema, generator, 0)
print(generator)
g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)
g_optim = optim.Adam(
generator.parameters(),
