def test_and_sample(cfg, model, device, test_loader, height, width, losses,
params, epoch):
test_loss = 0
model.eval()
with torch.no_grad():
for images, labels in test_loader:
images = images.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
normalized_images = images.float() / (cfg.color_levels - 1)
outputs = model(normalized_images, labels)
test_loss += F.cross_entropy(outputs, images, reduction='none')
test_loss = test_loss.mean().cpu() / len(test_loader.dataset)
wandb.log({"Test loss": test_loss})
print("Average test loss: {}".format(test_loss))
losses.append(test_loss)
params.append(model.state_dict())
samples = model.sample((3, height, width),
cfg.epoch_samples,
device=device)
save_samples(samples, TRAIN_SAMPLES_DIR,
'epoch{}_samples.png'.format(epoch + 1))
def _save_samples(self):
rows, columns = 5, 5
noise = np.random.normal(0, 1, (rows * columns, self._latent_dim))
generated_transactions = self._generator.predict(noise)
filenames = [
self._img_dir + ('/%07d.png' % self._epoch),
self._img_dir + '/last.png'
]
utils.save_samples(generated_transactions, rows, columns, filenames)
def train(data_root, model, total_epoch, batch_size, lrate):
X, Z, Lr = model.inputs()
d_loss, g_loss = model.loss(X, Z)
d_opt, g_opt = model.optimizer(d_loss, g_loss, Lr)
g_sample = model.sample(Z)
sample_size = batch_size
test_noise = utils.get_noise(sample_size, n_noise)
epoch_drop = 3
iterator, image_count = ImageIterator(data_root, batch_size,
model.image_size,
model.image_channels).get_iterator()
next_element = iterator.get_next()
total_batch = int(image_count / batch_size)
#learning_rate = lrate
#G_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='generator')
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
sess.run(iterator.initializer)
for epoch in range(total_epoch):
learning_rate = lrate * \
math.pow(0.2, math.floor((epoch + 1) / epoch_drop))
for step in range(total_batch):
batch_x = sess.run(next_element)
batch_z = utils.get_noise(batch_size, n_noise)
_, loss_val_D = sess.run([d_opt, d_loss],
feed_dict={
X: batch_x,
Z: batch_z,
Lr: learning_rate
})
_, loss_val_G = sess.run([g_opt, g_loss],
feed_dict={
Z: batch_z,
Lr: learning_rate
})
if step % 300 == 0:
#sample_size = 10
#noise = get_noise(sample_size, n_noise)
samples = sess.run(g_sample, feed_dict={Z: test_noise})
title = 'samples/%05d_%05d.png' % (epoch, step)
utils.save_samples(title, samples)
print('Epoch:', '%04d' % epoch,
'%05d/%05d' % (step, total_batch),
'D loss: {:.4}'.format(loss_val_D),
'G loss: {:.4}'.format(loss_val_G))
saver.save(sess, './models/dcgan', global_step=epoch)
def main():
parser = argparse.ArgumentParser(description='PixelCNN')
parser.add_argument('--causal-ksize', type=int, default=7,
help='Kernel size of causal convolution')
parser.add_argument('--hidden-ksize', type=int, default=7,
help='Kernel size of hidden layers convolutions')
parser.add_argument('--color-levels', type=int, default=2,
help='Number of levels to quantisize value of each channel of each pixel into')
parser.add_argument('--hidden-fmaps', type=int, default=30,
help='Number of feature maps in hidden layer')
parser.add_argument('--out-hidden-fmaps', type=int, default=10,
help='Number of feature maps in outer hidden layer')
parser.add_argument('--hidden-layers', type=int, default=6,
help='Number of layers of gated convolutions with mask of type "B"')
parser.add_argument('--cuda', type=str2bool, default=True,
help='Flag indicating whether CUDA should be used')
parser.add_argument('--model-path', '-m',
help="Path to model's saved parameters")
parser.add_argument('--output-fname', type=str, default='samples.png',
help='Name of output file (.png format)')
parser.add_argument('--label', '--l', type=int, default=-1,
help='Label of sampled images. -1 indicates random labels.')
parser.add_argument('--count', '-c', type=int, default=64,
help='Number of images to generate')
parser.add_argument('--height', type=int, default=28, help='Output image height')
parser.add_argument('--width', type=int, default=28, help='Output image width')
cfg = parser.parse_args()
OUTPUT_FILENAME = cfg.output_fname
model = PixelCNN(cfg=cfg)
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() and cfg.cuda else "cpu")
model.to(device)
model.load_state_dict(torch.load(cfg.model_path))
label = None if cfg.label == -1 else cfg.label
samples = model.sample((3, cfg.height, cfg.width), cfg.count, label=label, device=device)
save_samples(samples, OUTPUT_DIRNAME, OUTPUT_FILENAME)
def train(self):
def preprocess_func(x):
return tf.cast(x, tf.float32)/127.5 - 1.
logfile_path = os.path.join(self.results_dir, 'training_base_logfile.txt')
if not os.path.exists(logfile_path):
logfile = open(logfile_path, "w")
while True:
for dataset_path in sorted(os.listdir(self.data_dir)):
dataset = np.load(os.path.join(self.data_dir, dataset_path))
dataset = tf.data.Dataset.from_tensor_slices((dataset)).shuffle(self.batch_size*5).batch(self.batch_size, drop_remainder=True).map(preprocess_func)
dataset = self.strategy.experimental_distribute_dataset(dataset)
for x in dataset:
self.train_step(x)
i = self.i
if i%25 == 0: #updates the moving average every 25 steps , as it takes a very long time.
self._update_moving_avg()
if i%1000==0:
loss, gradnorm = to_numpy(self.train_loss.result()), to_numpy(self.train_gradnorm.result())
result_str = "Iteration: %d, Time Elapsed: %0.1f, Loss: %0.4f, gradnorm: %0.4f" % (i, time()-self.starttime, loss, gradnorm)
print(result_str)
logfile = open(logfile_path, "a")
logfile.write(result_str + "\n")
logfile.close()
self.train_loss.reset_states()
self.train_gradnorm.reset_states()
if i%10000==0:
self.ema_model.set_weights(self.moving_avg_weights)
self.save_objects(i//1000)
_, samples_ema = self.generate_samples(64)
savepath = self._getp('samples_ema_{}k.jpg'.format(i//1000), folder='results')
save_samples(samples_ema, savepath)
if i >= self.max_iterations:
print("Training complete. Recover your models from the {} folder and your training results from the {} folder".format(self.model_dir, self.results_dir))
print("Training was done on a maximum of {} iterations".format(i))
return True
self.i += 1
def test(self):
init = tf.initialize_all_variables()
with tf.Session() as sess:
sess.run(init)
self.saver.restore(sess, self.model_path)
sample_z = np.random.uniform(1,
-1,
size=[self.batch_size, self.z_dim])
sample_exec_time = sample_output(self.batch_size)
sample_feature = sess.run(self.fake_feature_vector,
feed_dict={
self.z: sample_z,
self.y: sample_exec_time
})
save_samples(sample_feature, sample_exec_time)
print("Test finish!")
def test(ckpt_root, model, batch_size):
X, Z, Lr, Kt = model.inputs()
g_sample = model.sample(Z, reuse=False)
sample_size = batch_size
test_noise = utils.get_uniform_noise(sample_size, n_noise)
saver = tf.train.Saver()
with tf.Session() as sess:
#sess.run(tf.global_variables_initializer())
saver.restore(sess, ckpt_root)
samples = sess.run(g_sample, feed_dict={Z: test_noise})
date = datetime.datetime.now()
title = 'samples/%s.png' % date
utils.save_samples(title, samples)
np.savetxt('samples/test_noise%s.txt' % date,
test_noise,
fmt='%2.5f',
delimiter=', ')
def sample(model, config, samples_dir, texture_path,
n_samples=20, n_z_samples=5, inverse=False):
utils.makedirs(samples_dir)
imgs = None
if inverse:
img_files = sorted(os.listdir(texture_path))[:n_samples]
img_files = [texture_path + file for file in img_files]
imgs = get_images(img_files)
imgs = [get_random_patch(img, config.npx) for img in imgs]
imgs = [np.reshape(img, (1,) + img.shape) for img in imgs]
all_samples = []
if inverse:
X = np.concatenate(imgs, axis=0)
global_noise = model.generate_z(X)
global_noise = model.generate_z_det(X)
z_samples = utils.sample_noise_tensor(config, n_z_samples, config.zx,
global_noise=global_noise,
per_each=True)
gen_samples = model.generate(z_samples)
gen_samples = model.generate_det(z_samples)
all_samples = [[img, list(gen_samples[n_z_samples*i:n_z_samples*(i+1)])]
for i, img in enumerate(imgs)]
all_samples = [np.concatenate(np.concatenate(samples, axis=0), axis=2)
for samples in all_samples]
all_samples = [np.concatenate(all_samples, axis=1)]
utils.save_samples(samples_dir, all_samples, ['inv_gens'])
all_samples = []
for i in range(n_samples):
global_noise = np.random.uniform(-1., 1., (1, config.nz_global, 1, 1))
z_samples = utils.sample_noise_tensor(config, n_z_samples, config.zx,
global_noise=global_noise)
gen_samples = model.generate_det(z_samples)
gen_samples = np.concatenate(gen_samples, axis=2)
all_samples.append(gen_samples)
all_samples = [np.concatenate(all_samples, axis=1)]
utils.save_samples(samples_dir, all_samples, ['gens'])
def test(ckpt_root, model, batch_size):
X, Z, Lr = model.inputs()
g_sample = model.sample(Z, reuse=False)
sample_size = batch_size
test_noise = utils.get_noise(sample_size, n_noise)
# TEST DCGAN interpolation
'''
test_noise[0] = np.array([[-0.01403, 0.56896, 1.41881, 0.02516, -1.36731, -0.92614, 1.17105, -0.17130, -0.11242, 0.35453, 0.06243, 0.41190, -0.18923, -1.10846, -0.10500, 0.65989, -0.19307, -0.32606, -1.77017, -0.38637, -0.82117, 0.53288, -0.38393, 1.16999, 0.02266, 0.36757, 0.13555, -1.06630, 0.00951, -0.04134, -0.29982, -0.83991, -0.04059, -0.56064, 0.39640, 0.29686, 0.42023, -1.15875, 0.19443, -0.89730, 0.37836, -2.48704, -0.03874, 0.04086, -0.35425, -0.02359, 0.56843, -0.45289, 1.79295, 0.98343, -0.99543, 0.70134, -1.43882, -0.10630, -0.39800, -1.90689, -0.16606, 0.01075, 0.11386, 0.08757, 0.25799, 1.06645, 0.07529, 1.17719, 1.38717, -0.93715, 0.60258, 0.64817, -0.70972, 1.49177, -0.58564, -1.47612, -0.49625, 2.30098, -0.08210, -0.22495, -0.47805, -0.72601, 0.58665, -0.63158, 0.04414, -0.05951, -0.92667, -0.07905, -2.26017, -0.29677, 0.93230, -0.06546, -0.46701, 1.49024, 0.01060, -0.86621, -0.65857, 0.42297, -1.43760, 0.53813, -0.13808, 0.23095, -0.78151, 0.63207
]])
test_noise[8] = np.array([[-0.30745, -1.80994, 0.84740, -0.01723, -0.25759, 1.62209, -0.01877, 1.31540, 1.80470, -1.76964, 2.06064, -0.62803, -0.94382, 0.85376, -0.26913, 0.69890, 1.52500, -0.62958, -0.97269, 1.81976, 1.46848, -0.10180, -0.14649, 0.82289, -0.21654, 0.63229, -0.61106, -0.84134, 0.95145, -0.84128, -0.02509, -0.14419, -0.46364, -0.00298, -0.23900, -1.37273, -1.16797, 1.10777, -1.56686, -0.60846, -0.18123, 1.95980, 0.58466, 0.64532, 1.01655, -1.00187, 0.07544, 0.31779, 1.55344, -0.41186, -0.14158, 0.07359, -1.02670, -0.14173, 0.10773, -0.64202, -1.56408, -1.96202, -0.13097, -0.05426, -2.26692, 0.04790, 0.03724, -0.55998, 0.11415, -1.97006, 0.56635, -1.29249, 0.32449, 0.37213, -0.77510, -0.09502, 2.44859, 0.68632, 0.48752, 0.18134, -1.14473, -0.09552, -0.62953, 0.28095, -1.04062, 0.39957, -1.39301, -0.29697, -0.99899, 1.91437, -1.94361, -0.38661, -0.04163, -0.09743, -0.87291, 1.00404, 0.51789, -0.78019, 1.43526, 0.16111, 1.26596, -0.12284, 0.74221, 1.53793
]])
test_noise[1] = (7*test_noise[0]+1*test_noise[8])/8
test_noise[2] = (6*test_noise[0]+2*test_noise[8])/8
test_noise[3] = (5*test_noise[0]+3*test_noise[8])/8
test_noise[4] = (4*test_noise[0]+4*test_noise[8])/8
test_noise[5] = (3*test_noise[0]+5*test_noise[8])/8
test_noise[6] = (2*test_noise[0]+6*test_noise[8])/8
test_noise[7] = (1*test_noise[0]+7*test_noise[8])/8
'''
saver = tf.train.Saver()
with tf.Session() as sess:
#sess.run(tf.global_variables_initializer())
saver.restore(sess, ckpt_root)
samples = sess.run(g_sample, feed_dict={Z: test_noise})
date = datetime.datetime.now()
title = 'samples/%s.png' % date
utils.save_samples(title, samples)
np.savetxt('samples/test_noise%s.txt' % date,
test_noise,
fmt='%2.5f',
delimiter=', ')
G_costs.append(G_cost_epoch_avg)
LOGGER.info(
"{} D_cost_train:{:.4f} | D_wass_train:{:.4f} | D_cost_valid:{:.4f} | D_wass_valid:{:.4f} | "
"G_cost:{:.4f}".format(time_since(start), D_cost_train_epoch_avg,
D_wass_train_epoch_avg, D_cost_valid_epoch_avg,
D_wass_valid_epoch_avg, G_cost_epoch_avg))
# Generate audio samples.
if epoch % epochs_per_sample == 0:
LOGGER.info("Generating samples...")
sample_out = netG(sample_noise_Var)
if cuda:
sample_out = sample_out.cpu()
sample_out = sample_out.data.numpy()
save_samples(sample_out, epoch, output_dir)
# TODO
# Early stopping by Inception Score(IS)
LOGGER.info('>>>>>>>Training finished !<<<<<<<')
# Save model
LOGGER.info("Saving models...")
netD_path = os.path.join(output_dir, "discriminator.pkl")
netG_path = os.path.join(output_dir, "generator.pkl")
torch.save(netD.state_dict(),
netD_path,
pickle_protocol=pickle.HIGHEST_PROTOCOL)
torch.save(netG.state_dict(),
netG_path,
def train():
# Fix Seed for Reproducibility #
torch.manual_seed(9)
if torch.cuda.is_available():
torch.cuda.manual_seed(9)
# Samples, Weights and Results Path #
paths = [config.samples_path, config.weights_path, config.plots_path]
paths = [make_dirs(path) for path in paths]
# Prepare Data Loader #
train_horse_loader, train_zebra_loader = get_horse2zebra_loader('train', config.batch_size)
val_horse_loader, val_zebra_loader = get_horse2zebra_loader('test', config.batch_size)
total_batch = min(len(train_horse_loader), len(train_zebra_loader))
# Image Pool #
masked_fake_A_pool = ImageMaskPool(config.pool_size)
masked_fake_B_pool = ImageMaskPool(config.pool_size)
# Prepare Networks #
Attn_A = Attention()
Attn_B = Attention()
G_A2B = Generator()
G_B2A = Generator()
D_A = Discriminator()
D_B = Discriminator()
networks = [Attn_A, Attn_B, G_A2B, G_B2A, D_A, D_B]
for network in networks:
network.to(device)
# Loss Function #
criterion_Adversarial = nn.MSELoss()
criterion_Cycle = nn.L1Loss()
# Optimizers #
D_optim = torch.optim.Adam(chain(D_A.parameters(), D_B.parameters()), lr=config.lr, betas=(0.5, 0.999))
G_optim = torch.optim.Adam(chain(Attn_A.parameters(), Attn_B.parameters(), G_A2B.parameters(), G_B2A.parameters()), lr=config.lr, betas=(0.5, 0.999))
D_optim_scheduler = get_lr_scheduler(D_optim)
G_optim_scheduler = get_lr_scheduler(G_optim)
# Lists #
D_A_losses, D_B_losses = [], []
G_A_losses, G_B_losses = [], []
# Train #
print("Training Unsupervised Attention-Guided GAN started with total epoch of {}.".format(config.num_epochs))
for epoch in range(config.num_epochs):
for i, (real_A, real_B) in enumerate(zip(train_horse_loader, train_zebra_loader)):
# Data Preparation #
real_A = real_A.to(device)
real_B = real_B.to(device)
# Initialize Optimizers #
D_optim.zero_grad()
G_optim.zero_grad()
###################
# Train Generator #
###################
set_requires_grad([D_A, D_B], requires_grad=False)
# Adversarial Loss using real A #
attn_A = Attn_A(real_A)
fake_B = G_A2B(real_A)
masked_fake_B = fake_B * attn_A + real_A * (1-attn_A)
masked_fake_B *= attn_A
prob_real_A = D_A(masked_fake_B)
real_labels = torch.ones(prob_real_A.size()).to(device)
G_loss_A = criterion_Adversarial(prob_real_A, real_labels)
# Adversarial Loss using real B #
attn_B = Attn_B(real_B)
fake_A = G_B2A(real_B)
masked_fake_A = fake_A * attn_B + real_B * (1-attn_B)
masked_fake_A *= attn_B
prob_real_B = D_B(masked_fake_A)
real_labels = torch.ones(prob_real_B.size()).to(device)
G_loss_B = criterion_Adversarial(prob_real_B, real_labels)
# Cycle Consistency Loss using real A #
attn_ABA = Attn_B(masked_fake_B)
fake_ABA = G_B2A(masked_fake_B)
masked_fake_ABA = fake_ABA * attn_ABA + masked_fake_B * (1 - attn_ABA)
# Cycle Consistency Loss using real B #
attn_BAB = Attn_A(masked_fake_A)
fake_BAB = G_A2B(masked_fake_A)
masked_fake_BAB = fake_BAB * attn_BAB + masked_fake_A * (1 - attn_BAB)
# Cycle Consistency Loss #
G_cycle_loss_A = config.lambda_cycle * criterion_Cycle(masked_fake_ABA, real_A)
G_cycle_loss_B = config.lambda_cycle * criterion_Cycle(masked_fake_BAB, real_B)
# Total Generator Loss #
G_loss = G_loss_A + G_loss_B + G_cycle_loss_A + G_cycle_loss_B
# Back Propagation and Update #
G_loss.backward()
G_optim.step()
#######################
# Train Discriminator #
#######################
set_requires_grad([D_A, D_B], requires_grad=True)
# Train Discriminator A using real A #
prob_real_A = D_A(real_B)
real_labels = torch.ones(prob_real_A.size()).to(device)
D_loss_real_A = criterion_Adversarial(prob_real_A, real_labels)
# Add Pooling #
masked_fake_B, attn_A = masked_fake_B_pool.query(masked_fake_B, attn_A)
masked_fake_B *= attn_A
# Train Discriminator A using fake B #
prob_fake_B = D_A(masked_fake_B.detach())
fake_labels = torch.zeros(prob_fake_B.size()).to(device)
D_loss_fake_A = criterion_Adversarial(prob_fake_B, fake_labels)
D_loss_A = (D_loss_real_A + D_loss_fake_A).mean()
# Train Discriminator B using real B #
prob_real_B = D_B(real_A)
real_labels = torch.ones(prob_real_B.size()).to(device)
D_loss_real_B = criterion_Adversarial(prob_real_B, real_labels)
# Add Pooling #
masked_fake_A, attn_B = masked_fake_A_pool.query(masked_fake_A, attn_B)
masked_fake_A *= attn_B
# Train Discriminator B using fake A #
prob_fake_A = D_B(masked_fake_A.detach())
fake_labels = torch.zeros(prob_fake_A.size()).to(device)
D_loss_fake_B = criterion_Adversarial(prob_fake_A, fake_labels)
D_loss_B = (D_loss_real_B + D_loss_fake_B).mean()
# Calculate Total Discriminator Loss #
D_loss = D_loss_A + D_loss_B
# Back Propagation and Update #
D_loss.backward()
D_optim.step()
# Add items to Lists #
D_A_losses.append(D_loss_A.item())
D_B_losses.append(D_loss_B.item())
G_A_losses.append(G_loss_A.item())
G_B_losses.append(G_loss_B.item())
####################
# Print Statistics #
####################
if (i+1) % config.print_every == 0:
print("UAG-GAN | Epoch [{}/{}] | Iteration [{}/{}] | D A Losses {:.4f} | D B Losses {:.4f} | G A Losses {:.4f} | G B Losses {:.4f}".
format(epoch+1, config.num_epochs, i+1, total_batch, np.average(D_A_losses), np.average(D_B_losses), np.average(G_A_losses), np.average(G_B_losses)))
# Save Sample Images #
save_samples(val_horse_loader, val_zebra_loader, G_A2B, G_B2A, Attn_A, Attn_B, epoch, config.samples_path)
# Adjust Learning Rate #
D_optim_scheduler.step()
G_optim_scheduler.step()
# Save Model Weights #
if (epoch + 1) % config.save_every == 0:
torch.save(G_A2B.state_dict(), os.path.join(config.weights_path, 'UAG-GAN_Generator_A2B_Epoch_{}.pkl'.format(epoch+1)))
torch.save(G_B2A.state_dict(), os.path.join(config.weights_path, 'UAG-GAN_Generator_B2A_Epoch_{}.pkl'.format(epoch+1)))
torch.save(Attn_A.state_dict(), os.path.join(config.weights_path, 'UAG-GAN_Attention_A_Epoch_{}.pkl'.format(epoch+1)))
torch.save(Attn_B.state_dict(), os.path.join(config.weights_path, 'UAG-GAN_Attention_B_Epoch_{}.pkl'.format(epoch+1)))
# Make a GIF file #
make_gifs_train("UAG-GAN", config.samples_path)
# Plot Losses #
plot_losses(D_A_losses, D_B_losses, G_A_losses, G_B_losses, config.num_epochs, config.plots_path)
print("Training finished.")
def train():
# Fix Seed for Reproducibility #
torch.manual_seed(9)
if torch.cuda.is_available():
torch.cuda.manual_seed(9)
# Samples, Weights, and Plots Path #
paths = [config.samples_path, config.weights_path, config.plots_path]
paths = [make_dirs(path) for path in paths]
# Prepare Data Loader #
train_loader_selfie, train_loader_anime = get_selfie2anime_loader(
'train', config.batch_size)
total_batch = max(len(train_loader_selfie), len(train_loader_anime))
test_loader_selfie, test_loader_anime = get_selfie2anime_loader(
'test', config.val_batch_size)
# Prepare Networks #
D_A = Discriminator(num_layers=7)
D_B = Discriminator(num_layers=7)
L_A = Discriminator(num_layers=5)
L_B = Discriminator(num_layers=5)
G_A2B = Generator(image_size=config.crop_size,
num_blocks=config.num_blocks)
G_B2A = Generator(image_size=config.crop_size,
num_blocks=config.num_blocks)
networks = [D_A, D_B, L_A, L_B, G_A2B, G_B2A]
for network in networks:
network.to(device)
# Loss Function #
Adversarial_loss = nn.MSELoss()
Cycle_loss = nn.L1Loss()
BCE_loss = nn.BCEWithLogitsLoss()
# Optimizers #
D_optim = torch.optim.Adam(chain(D_A.parameters(), D_B.parameters(),
L_A.parameters(), L_B.parameters()),
lr=config.lr,
betas=(0.5, 0.999),
weight_decay=0.0001)
G_optim = torch.optim.Adam(chain(G_A2B.parameters(), G_B2A.parameters()),
lr=config.lr,
betas=(0.5, 0.999),
weight_decay=0.0001)
D_optim_scheduler = get_lr_scheduler(D_optim)
G_optim_scheduler = get_lr_scheduler(G_optim)
# Rho Clipper to constraint the value of rho in AdaILN and ILN #
Rho_Clipper = RhoClipper(0, 1)
# Lists #
D_losses = []
G_losses = []
# Train #
print("Training U-GAT-IT started with total epoch of {}.".format(
config.num_epochs))
for epoch in range(config.num_epochs):
for i, (selfie, anime) in enumerate(
zip(train_loader_selfie, train_loader_anime)):
# Data Preparation #
real_A = selfie.to(device)
real_B = anime.to(device)
# Initialize Optimizers #
D_optim.zero_grad()
G_optim.zero_grad()
#######################
# Train Discriminator #
#######################
set_requires_grad([D_A, D_B, L_A, L_B], requires_grad=True)
# Forward Data #
fake_B, _, _ = G_A2B(real_A)
fake_A, _, _ = G_B2A(real_B)
G_real_A, G_real_A_cam, _ = D_A(real_A)
L_real_A, L_real_A_cam, _ = L_A(real_A)
G_real_B, G_real_B_cam, _ = D_B(real_B)
L_real_B, L_real_B_cam, _ = L_B(real_B)
G_fake_A, G_fake_A_cam, _ = D_A(fake_A)
L_fake_A, L_fake_A_cam, _ = L_A(fake_A)
G_fake_B, G_fake_B_cam, _ = D_B(fake_B)
L_fake_B, L_fake_B_cam, _ = L_B(fake_B)
# Adversarial Loss of Discriminator #
real_labels = torch.ones(G_real_A.shape).to(device)
D_ad_real_loss_GA = Adversarial_loss(G_real_A, real_labels)
fake_labels = torch.zeros(G_fake_A.shape).to(device)
D_ad_fake_loss_GA = Adversarial_loss(G_fake_A, fake_labels)
D_ad_loss_GA = D_ad_real_loss_GA + D_ad_fake_loss_GA
real_labels = torch.ones(G_real_A_cam.shape).to(device)
D_ad_cam_real_loss_GA = Adversarial_loss(G_real_A_cam, real_labels)
fake_labels = torch.zeros(G_fake_A_cam.shape).to(device)
D_ad_cam_fake_loss_GA = Adversarial_loss(G_fake_A_cam, fake_labels)
D_ad_cam_loss_GA = D_ad_cam_real_loss_GA + D_ad_cam_fake_loss_GA
real_labels = torch.ones(G_real_B.shape).to(device)
D_ad_real_loss_GB = Adversarial_loss(G_real_B, real_labels)
fake_labels = torch.zeros(G_fake_B.shape).to(device)
D_ad_fake_loss_GB = Adversarial_loss(G_fake_B, fake_labels)
D_ad_loss_GB = D_ad_real_loss_GB + D_ad_fake_loss_GB
real_labels = torch.ones(G_real_B_cam.shape).to(device)
D_ad_cam_real_loss_GB = Adversarial_loss(G_real_B_cam, real_labels)
fake_labels = torch.zeros(G_fake_B_cam.shape).to(device)
D_ad_cam_fake_loss_GB = Adversarial_loss(G_fake_B_cam, fake_labels)
D_ad_cam_loss_GB = D_ad_cam_real_loss_GB + D_ad_cam_fake_loss_GB
# Adversarial Loss of L #
real_labels = torch.ones(L_real_A.shape).to(device)
D_ad_real_loss_LA = Adversarial_loss(L_real_A, real_labels)
fake_labels = torch.zeros(L_fake_A.shape).to(device)
D_ad_fake_loss_LA = Adversarial_loss(L_fake_A, fake_labels)
D_ad_loss_LA = D_ad_real_loss_LA + D_ad_fake_loss_LA
real_labels = torch.ones(L_real_A_cam.shape).to(device)
D_ad_cam_real_loss_LA = Adversarial_loss(L_real_A_cam, real_labels)
fake_labels = torch.zeros(L_fake_A_cam.shape).to(device)
D_ad_cam_fake_loss_LA = Adversarial_loss(L_fake_A_cam, fake_labels)
D_ad_cam_loss_LA = D_ad_cam_real_loss_LA + D_ad_cam_fake_loss_LA
real_labels = torch.ones(L_real_B.shape).to(device)
D_ad_real_loss_LB = Adversarial_loss(L_real_B, real_labels)
fake_labels = torch.zeros(L_fake_B.shape).to(device)
D_ad_fake_loss_LB = Adversarial_loss(L_fake_B, fake_labels)
D_ad_loss_LB = D_ad_real_loss_LB + D_ad_fake_loss_LB
real_labels = torch.ones(L_real_B_cam.shape).to(device)
D_ad_cam_real_loss_LB = Adversarial_loss(L_real_B_cam, real_labels)
fake_labels = torch.zeros(L_fake_B_cam.shape).to(device)
D_ad_cam_fake_loss_LB = Adversarial_loss(L_fake_B_cam, fake_labels)
D_ad_cam_loss_LB = D_ad_cam_real_loss_LB + D_ad_cam_fake_loss_LB
# Calculate Each Discriminator Loss #
D_loss_A = D_ad_loss_GA + D_ad_cam_loss_GA + D_ad_loss_LA + D_ad_cam_loss_LA
D_loss_B = D_ad_loss_GB + D_ad_cam_loss_GB + D_ad_loss_LB + D_ad_cam_loss_LB
# Calculate Total Discriminator Loss #
D_loss = D_loss_A + D_loss_B
# Back Propagation and Update #
D_loss.backward()
D_optim.step()
###################
# Train Generator #
###################
set_requires_grad([D_A, D_B, L_A, L_B], requires_grad=False)
# Forward Data #
fake_B, fake_B_cam, _ = G_A2B(real_A)
fake_A, fake_A_cam, _ = G_B2A(real_B)
fake_ABA, _, _ = G_B2A(fake_B)
fake_BAB, _, _ = G_A2B(fake_A)
fake_A2A, fake_A2A_cam, _ = G_A2B(real_A)
fake_B2B, fake_B2B_cam, _ = G_B2A(real_B)
G_fake_A, G_fake_A_cam, _ = D_A(fake_A)
L_fake_A, L_fake_A_cam, _ = L_A(fake_A)
G_fake_B, G_fake_B_cam, _ = D_B(fake_B)
L_fake_B, L_fake_B_cam, _ = L_B(fake_B)
# Adversarial Loss of Generator #
real_labels = torch.ones(G_fake_A.shape).to(device)
G_adv_fake_loss_A = Adversarial_loss(G_fake_A, real_labels)
real_labels = torch.ones(G_fake_A_cam.shape).to(device)
G_adv_cam_fake_loss_A = Adversarial_loss(G_fake_A_cam, real_labels)
G_adv_loss_A = G_adv_fake_loss_A + G_adv_cam_fake_loss_A
real_labels = torch.ones(G_fake_B.shape).to(device)
G_adv_fake_loss_B = Adversarial_loss(G_fake_B, real_labels)
real_labels = torch.ones(G_fake_B_cam.shape).to(device)
G_adv_cam_fake_loss_B = Adversarial_loss(G_fake_B_cam, real_labels)
G_adv_loss_B = G_adv_fake_loss_B + G_adv_cam_fake_loss_B
# Adversarial Loss of L #
real_labels = torch.ones(L_fake_A.shape).to(device)
L_adv_fake_loss_A = Adversarial_loss(L_fake_A, real_labels)
real_labels = torch.ones(L_fake_A_cam.shape).to(device)
L_adv_cam_fake_loss_A = Adversarial_loss(L_fake_A_cam, real_labels)
L_adv_loss_A = L_adv_fake_loss_A + L_adv_cam_fake_loss_A
real_labels = torch.ones(L_fake_B.shape).to(device)
L_adv_fake_loss_B = Adversarial_loss(L_fake_B, real_labels)
real_labels = torch.ones(L_fake_B_cam.shape).to(device)
L_adv_cam_fake_loss_B = Adversarial_loss(L_fake_B_cam, real_labels)
L_adv_loss_B = L_adv_fake_loss_B + L_adv_cam_fake_loss_B
# Cycle Consistency Loss #
G_recon_loss_A = Cycle_loss(fake_ABA, real_A)
G_recon_loss_B = Cycle_loss(fake_BAB, real_B)
G_identity_loss_A = Cycle_loss(fake_A2A, real_A)
G_identity_loss_B = Cycle_loss(fake_B2B, real_B)
G_cycle_loss_A = G_recon_loss_A + G_identity_loss_A
G_cycle_loss_B = G_recon_loss_B + G_identity_loss_B
# CAM Loss #
real_labels = torch.ones(fake_A_cam.shape).to(device)
G_cam_real_loss_A = BCE_loss(fake_A_cam, real_labels)
fake_labels = torch.zeros(fake_A2A_cam.shape).to(device)
G_cam_fake_loss_A = BCE_loss(fake_A2A_cam, fake_labels)
G_cam_loss_A = G_cam_real_loss_A + G_cam_fake_loss_A
real_labels = torch.ones(fake_B_cam.shape).to(device)
G_cam_real_loss_B = BCE_loss(fake_B_cam, real_labels)
fake_labels = torch.zeros(fake_B2B_cam.shape).to(device)
G_cam_fake_loss_B = BCE_loss(fake_B2B_cam, fake_labels)
G_cam_loss_B = G_cam_real_loss_B + G_cam_fake_loss_B
# Calculate Each Generator Loss #
G_loss_A = G_adv_loss_A + L_adv_loss_A + config.lambda_cycle * G_cycle_loss_A + config.lambda_cam * G_cam_loss_A
G_loss_B = G_adv_loss_B + L_adv_loss_B + config.lambda_cycle * G_cycle_loss_B + config.lambda_cam * G_cam_loss_B
# Calculate Total Generator Loss #
G_loss = G_loss_A + G_loss_B
# Back Propagation and Update #
G_loss.backward()
G_optim.step()
# Apply Rho Clipper to Generators #
G_A2B.apply(Rho_Clipper)
G_B2A.apply(Rho_Clipper)
# Add items to Lists #
D_losses.append(D_loss.item())
G_losses.append(G_loss.item())
####################
# Print Statistics #
####################
if (i + 1) % config.print_every == 0:
print(
"U-GAT-IT | Epochs [{}/{}] | Iterations [{}/{}] | D Loss {:.4f} | G Loss {:.4f}"
.format(epoch + 1, config.num_epochs, i + 1, total_batch,
np.average(D_losses), np.average(G_losses)))
# Save Sample Images #
save_samples(test_loader_selfie, G_A2B, epoch,
config.samples_path)
# Adjust Learning Rate #
D_optim_scheduler.step()
G_optim_scheduler.step()
# Save Model Weights #
if (epoch + 1) % config.save_every == 0:
torch.save(
D_A.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_D_A_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
D_B.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_D_B_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
L_A.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_L_A_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
L_B.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_L_B_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
G_A2B.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_G_A2B_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
G_B2A.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_G_B2A_Epoch_{}.pkl'.format(epoch + 1)))
# Plot Losses #
plot_losses(D_losses, G_losses, config.num_epochs, config.plots_path)
# Make a GIF file #
make_gifs_train('U-GAT-IT', config.samples_path)
print("Training finished.")
# raw_data.append(spikes)
# Reshape to required format
data = data.reshape((1, imageSize, imageSize))
recv_data[rank - 1] = data
comm.Send(recv_data, dest=0)
# print('done generating data, in sec: {}'.format(time.time() - t))
if rank == 0:
for i in range(1, size):
comm.Recv(recv_data, source=i)
binned_data[i - 1] = recv_data[i - 1]
if comm.rank == 0:
normed_data = np.empty((num_samples, 1, imageSize, imageSize),
dtype=np.float32)
for i, bst in enumerate(binned_data):
# Normalize data
normed_data[i] = np.divide(bst, np.max(bst))
# data = (data - data.mean()) / data.std()
save_dict['binned_data'] = binned_data
save_dict['normed_data'] = normed_data
save_dict['num_samples'] = num_samples
save_dict['imageSize'] = imageSize
# save_dict['spikes'] = raw_data
save_dict['data_type'] = data_type
utils.save_samples(save_dict,
path='.',
filename='data_NS{}_IS{}_type-{}.npy'.format(
num_samples, imageSize, data_type))
loss += value
loss.backward()
optimizer.step()
print("[Epoch %d/%d] [Batch %d/%d] [loss: %f] " %
(epoch, args.epoch, idx, len(train_loader), loss))
if epoch % args.val_epoch == 0:
mse, ssim = utils.evaluate(test_loader, unet, args)
print("mse = {}, ssim = {}".format(mse, ssim))
score = 1 - mse / 100.0 + ssim
if score > best_score:
# if ssim > best_ssim:
torch.save(
unet.state_dict(),
os.path.join(args.save_dir, 'unet_{}.pth'.format(epoch)))
best_score = score
# best_ssim = ssim
utils.save_samples(test_loader, unet, save_samples_dir, epoch)
print('ssim = {}, mse = {}, score = {}, best_score = {} '.format(
ssim, mse, score, best_score))
print('best_score = {} '.format(best_score))
'''
Reference:
https://github.com/ceshine/fast-neural-style/blob/master/notebooks/01-image-style-transfer.ipynb
'''
def test(self, training=False):
# input images for the model to train on
x = tf.placeholder(tf.float32,
shape=(None, self.height, self.width,
self.channels),
name='inputs')
# correct images for the model to check against when learning
if (self.config == '--MNIST'):
y = tf.placeholder(tf.float32,
shape=(None, self.height, self.width,
self.channels),
name='correct_images')
elif (self.config == '--CIFAR' or self.config == '--FREY'):
y = tf.placeholder(tf.float32,
shape=(None, self.height, self.width,
256 * self.channels),
name='correct_images')
# model is in the training phase
is_training = tf.placeholder(tf.bool, name='training')
# build out the network architecture
self.build_network(x, y, is_training, 'test')
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, '/tmp/model.ckpt')
loss = []
NLL = []
for i in range(self.test_inputs.shape[0]):
if (self.config == '--MNIST'):
binaryImage = binarize(self.test_inputs[i])
binaryImage = np.reshape(binaryImage, (1, 28, 28, 1))
loss.append(
sess.run(self.loss,
feed_dict={
x: binaryImage,
y: binaryImage,
is_training: training
}))
NLL.append(
sess.run(self.NLL,
feed_dict={
x: binaryImage,
y: binaryImage,
is_training: training
}))
elif (self.config == '--FREY'):
image = self.test_inputs[i]
label = sess.run(
tf.one_hot(np.reshape(image, (-1, 28, 20)),
256,
axis=-1))
image = 2. * (np.reshape(
image, (1, 28, 20, 1)).astype('float32') / 255.) - 1.
loss.append(
sess.run(self.loss,
feed_dict={
x: image,
y: label,
is_training: training
}))
NLL.append(
sess.run(self.NLL,
feed_dict={
x: image,
y: label,
is_training: training
}))
loss = np.mean(loss)
NLL = np.mean(NLL)
print("Test loss: ", loss, "NLL: ", NLL)
#save_samples(self.test_inputs[:22], self.height, self.width)
# remove bottom half from images
images = trim_images(self.test_inputs)
if (self.config == '--MNIST'):
images = images[:22]
self.test_inputs = self.test_inputs[:22]
elif (self.config == '--CIFAR' or self.config == '--FREY'):
images = 2 * (images[:22].astype('float32') / 255.) - 1.
self.test_inputs = sess.run(
tf.one_hot(np.reshape(self.test_inputs[:22], (-1, 28, 20)),
256,
axis=-1))
#save_samples(images, self.height, self.width)
# use model to generate bottom half of images
for i in range(self.height // 2, self.height):
for j in range(self.width):
for k in range(self.channels):
probs = sess.run(self.pred,
feed_dict={
x: images,
y: self.test_inputs,
is_training: training
})
if (self.config == '--MNIST'):
sample = binarize(probs)
images[:, i, j, k] = sample[:, i, j, k]
elif (self.config == '--FREY'):
sample = sample_categorical(probs[:, i, j])
images[:, i, j,
k] = 2 * (sample.astype('float32') /
255.) - 1.
save_samples(images, self.height, self.width)
'%s/fake_samples_epoch_%03d.png' %
(opt.outf, epoch),
normalize=False)
'''
1-dim: list with all the data, listed according the epochs
2-dim: list containing lists of integer and tuple,
integer indicates the epoch, the tuple contains the step index
and the output data,
[int, tuple]
3-dim: tuple of integer and data as torch.FloatTensor, the integer
indicates the step index of the corresponding batch in the loop
(int, FloatTensor)
4-dim: FloatTensor of shape 64x1x64x64:
5-dim:
64 samples x
Channel number (here always only 1) x
64x64 normalized binned data
'''
save_dict['fake_data'][epoch].append((i, fake.data))
# do checkpointing
checkpoint_path = os.path.join(opt.outf, 'checkpoints')
if not os.path.exists(checkpoint_path):
os.mkdir(checkpoint_path)
torch.save(netG.state_dict(),
'%s/netG_epoch_%d.pth' % (checkpoint_path, epoch))
torch.save(netD.state_dict(),
'%s/netD_epoch_%d.pth' % (checkpoint_path, epoch))
utils.save_samples(save_dict, path=opt.outf, filename='results.npy')
writer.close()
def train(self):
def preprocess_func(x, y):
x = tf.cast(x, tf.float32)
y = tf.cast(y, tf.float32) / 127.5 - 1.
return x, y
logfile_path = os.path.join(self.results_dir,
'training_stg2_logfile.txt')
if not os.path.exists(logfile_path):
logfile = open(logfile_path, "w")
lsdir_x = [f for f in os.listdir(self.stg2_data_dir) if 'data_x' in f]
lsdir_y = [f for f in os.listdir(self.stg2_data_dir) if 'data_y' in f]
lsdir_x, lsdir_y = sorted(lsdir_x), sorted(lsdir_y)
assert len(lsdir_x) == len(
lsdir_y
), "You are missing an x or y dataset shard. You may have deleted it."
while True:
for xtr_path, ytr_path in zip(lsdir_x, lsdir_y):
x_tr = np.load(os.path.join(self.stg2_data_dir, xtr_path))
y_tr = np.load(os.path.join(self.stg2_data_dir, ytr_path))
dataset = tf.data.Dataset.from_tensor_slices(
(x_tr, y_tr)).shuffle(self.batch_size * 5).batch(
self.batch_size,
drop_remainder=True).map(preprocess_func)
dataset = self.strategy.experimental_distribute_dataset(
dataset)
for x, y in dataset:
self.train_step(x, y)
i = self.i
if i % 25 == 0: #updates the moving average every 25 steps , as it takes a very long time.
self._update_moving_avg()
if i % 1000 == 0:
loss, gradnorm = to_numpy(
self.train_loss.result()), to_numpy(
self.train_gradnorm.result())
result_str = "Iteration: %d, Time Elapsed: %0.1f, Loss: %0.4f, gradnorm: %0.4f" % (
i, time() - self.starttime, loss, gradnorm)
print(result_str)
logfile = open(logfile_path, "a")
logfile.write(result_str + "\n")
logfile.close()
self.train_loss.reset_states()
self.train_gradnorm.reset_states()
if i % 10000 == 0:
self.ema_model.set_weights(self.moving_avg_weights)
self.save_objects(i // 1000, stg2=True)
_, samples_ema = self.generate_samples(64)
savepath = self._getp('samples_ema_{}k.jpg'.format(
i // 1000),
folder='results')
save_samples(samples_ema, savepath)
if i >= self.max_iterations:
print(
"Training complete. Recover your models from the {} folder and your training results from the {} folder"
.format(self.model_dir, self.results_dir))
print(
"Training was done on a maximum of {} iterations".
format(i))
return True
self.i += 1
writer.add_scalar("loss/penalty", penalty, iters)
writer.add_scalar("loss/mi", loss_mi, iters)
if iters % args.log_interval == 0:
print(
"Train Iter: {}/{} ({:.0f}%)\t"
"D_costs: {} G_costs: {} Time: {:5.3f}".format(
iters,
args.iters,
(args.log_interval * iters) / args.iters,
np.asarray(e_costs).mean(0),
np.asarray(g_costs).mean(0),
(time.time() - start_time) / args.log_interval,
)
)
img = save_samples(netG, args)
writer.add_image("samples/generated", img, iters)
e_costs = []
g_costs = []
start_time = time.time()
if iters % args.save_interval == 0:
netG.eval()
print("-" * 100)
n_modes, kld = mc_eval.count_modes(netG)
print("-" * 100)
netG.train()
writer.add_scalar("metrics/mode_count", n_modes, iters)
writer.add_scalar("metrics/kl_divergence", kld, iters)
def train(model, config, logger, options, model_dir, samples_dir,
inverse=False, save_step=10, n_samples=20):
utils.makedirs(samples_dir)
losses = defaultdict(list)
for epoch in tqdm(range(options.n_epochs), file=sys.stdout):
logger.info("Epoch {}".format(epoch))
samples_generator = config.data_iter(
options.data, options.b_size, inverse=inverse, n_samples=n_samples)
entropy_epoch = []
for it in tqdm(range(options.n_iters), file=sys.stdout):
Z_global = None
if inverse:
Z_global = np.random.uniform(
-1., 1., (options.b_size, config.nz_global, 1, 1))
Z_samples = utils.sample_noise_tensor(
config, options.b_size, config.zx, global_noise=Z_global)
X_samples = next(samples_generator)
if it % (config.k + 1) != 0:
if inverse == 0:
losses['G_iter'].append(model.train_g(Z_samples))
elif inverse == 1:
losses['G_iter'].append(model.train_g(X_samples, Z_samples,
Z_global))
elif inverse >= 2:
# losses['G_iter'].append(model.train_g(X_samples[0],
# Z_samples))
G_loss, entropy = model.train_g(X_samples[0], Z_samples)
losses['G_iter'].append(G_loss)
entropy_epoch.append(entropy)
else:
if inverse == 0:
losses['D_iter'].append(model.train_d(X_samples, Z_samples))
elif inverse == 1:
losses['D_iter'].append(model.train_d(X_samples, Z_samples,
Z_global))
elif inverse >= 2:
losses['D_iter'].append(model.train_d(
X_samples[0], X_samples[1], Z_samples))
msg = "Gloss = {}, Dloss = {}"
msg += "\n e_min = {}, e_max = {}, e_mean = {}, e_med = {}"
e_min, e_max, e_mean, e_med = [f(entropy_epoch) for f in
[np.min, np.max, np.mean, np.median]]
losses['G_epoch'].append(np.mean(losses['G_iter'][-options.n_iters:]))
losses['D_epoch'].append(np.mean(losses['D_iter'][-options.n_iters:]))
logger.info(msg.format(losses['G_epoch'][-1], losses['D_epoch'][-1],
e_min, e_max, e_mean, e_med))
X = next(samples_generator)
real_samples, gen_samples, large_sample = utils.sample_after_iteration(
model, X, inverse, config, options.b_size)
utils.save_samples(
samples_dir, [real_samples, gen_samples, large_sample],
['real', 'gen', 'large'], epoch=epoch)
utils.save_plots(samples_dir, losses, epoch, options.n_iters)
if (epoch+1) % save_step == 0:
model_file = 'epoch_{:04d}.model'.format(epoch)
model.save(os.path.join(model_dir, model_file))
def train(data_root, model, total_epoch, batch_size, lrate):
X, Z, Lr, Kt = model.inputs()
d_loss, g_loss, real_loss, fake_loss = model.loss(X, Z, Kt)
d_opt, g_opt = model.optimizer(d_loss, g_loss, Lr)
g_sample = model.sample(Z)
sample_size = batch_size
test_noise = utils.get_noise(sample_size, n_noise)
epoch_drop = 3
_lambda = 0.001
_gamma = 0.5
_kt = 0.0
iterator, image_count = ImageIterator(data_root, batch_size,
model.image_size,
model.image_channels).get_iterator()
next_element = iterator.get_next()
measure = real_loss + tf.abs(_gamma * real_loss - fake_loss)
tf.summary.scalar('measure', measure)
merged = tf.summary.merge_all()
total_batch = int(image_count / batch_size)
#learning_rate = lrate
#G_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='generator')
saver = tf.train.Saver()
with tf.Session() as sess:
train_writer = tf.summary.FileWriter('./logs', sess.graph)
sess.run(tf.global_variables_initializer())
sess.run(iterator.initializer)
for epoch in range(total_epoch):
learning_rate = lrate * \
math.pow(0.2, math.floor((epoch + 1) / epoch_drop))
for step in range(total_batch):
batch_x = sess.run(next_element)
batch_z = utils.get_uniform_noise(batch_size, n_noise)
_, val_real_loss = sess.run([d_opt, real_loss],
feed_dict={
X: batch_x,
Z: batch_z,
Lr: learning_rate,
Kt: _kt
})
_, val_fake_loss = sess.run([g_opt, fake_loss],
feed_dict={
Z: batch_z,
Lr: learning_rate,
Kt: _kt
})
_kt = _kt + _lambda * (_gamma * val_real_loss - val_fake_loss)
if step % 300 == 0:
summary = sess.run(merged,
feed_dict={
X: batch_x,
Z: batch_z,
Lr: learning_rate,
Kt: _kt
})
train_writer.add_summary(summary,
epoch * total_batch + step)
val_measure = val_real_loss + np.abs(
_gamma * val_real_loss - val_fake_loss)
print('Epoch:', '%04d' % epoch,
'%05d/%05d' % (step, total_batch),
'measure: {:.4}'.format(val_measure))
#sample_size = 10
#noise = get_noise(sample_size, n_noise)
samples = sess.run(g_sample, feed_dict={Z: test_noise})
title = 'samples/%05d_%05d.png' % (epoch, step)
utils.save_samples(title, samples)
saver.save(sess, './models/began', global_step=epoch)
type=str,
default='training_results/0000-00-00_00-00-00',
help='model save path.')
parser.add_argument('--save_latent',
action='store_true',
default=False,
help='save latent variables')
options = parser.parse_args()
print(options)
if __name__ == "__main__":
test_info, test_data, _ = load_samples(options.data_path)
model = load_model(options.model_path)
predictions = model.predict(test_data)
save_samples(os.path.join(options.save_path, 'recons'), predictions,
test_info)
#print(test_data.shape)
save_recon_error(options.save_path, test_data, predictions, test_info)
if options.save_latent:
encoder = Model(inputs=model.get_layer('sequential_1').model.input,
outputs=model.get_layer('sequential_1').model.output)
latents = encoder.predict(test_data)
save_latent_variables(os.path.join(options.save_path, 'latents'),
latents, test_info)
# ()()
# ('') HAANJU.YOO
def train_wgan(model_gen,
model_dis,
train_gen,
valid_data,
test_data,
num_epochs,
batches_per_epoch,
batch_size,
output_dir=None,
lmbda=0.1,
use_cuda=True,
discriminator_updates=5,
epochs_per_sample=10,
sample_size=20,
lr=1e-4,
beta_1=0.5,
beta_2=0.9,
latent_dim=100):
if use_cuda:
model_gen = model_gen.cuda()
model_dis = model_dis.cuda()
# Initialize optimizers for each model
optimizer_gen = optim.Adam(model_gen.parameters(),
lr=lr,
betas=(beta_1, beta_2))
optimizer_dis = optim.Adam(model_dis.parameters(),
lr=lr,
betas=(beta_1, beta_2))
# Sample noise used for seeing the evolution of generated output samples throughout training
sample_noise = torch.Tensor(sample_size, latent_dim).uniform_(-1, 1)
if use_cuda:
sample_noise = sample_noise.cuda()
sample_noise_v = autograd.Variable(sample_noise)
samples = {}
history = []
train_iter = iter(train_gen)
valid_data_v = np_to_input_var(valid_data['X'], use_cuda)
test_data_v = np_to_input_var(test_data['X'], use_cuda)
# Loop over the dataset multiple times
for epoch in range(num_epochs):
LOGGER.info("Epoch: {}/{}".format(epoch + 1, num_epochs))
epoch_history = []
for batch_idx in range(batches_per_epoch):
# Set model parameters to require gradients to be computed and stored
for p in model_dis.parameters():
p.requires_grad = True
# Initialize the metrics for this batch
batch_history = {'discriminator': [], 'generator': {}}
# Discriminator Training Phase:
# -> Train discriminator k times
for iter_d in range(discriminator_updates):
# Get real examples
real_data_v = np_to_input_var(next(train_iter)['X'], use_cuda)
# Get noise
noise = torch.Tensor(batch_size, latent_dim).uniform_(-1, 1)
if use_cuda:
noise = noise.cuda()
noise_v = autograd.Variable(
noise, volatile=True) # totally freeze model_gen
# Get new batch of real training data
D_cost_train, D_wass_train = compute_discr_loss_terms(
model_dis,
model_gen,
real_data_v,
noise_v,
batch_size,
latent_dim,
lmbda,
use_cuda,
compute_grads=True)
# Update the discriminator
optimizer_dis.step()
D_cost_valid, D_wass_valid = compute_discr_loss_terms(
model_dis,
model_gen,
valid_data_v,
noise_v,
batch_size,
latent_dim,
lmbda,
use_cuda,
compute_grads=False)
if use_cuda:
D_cost_train = D_cost_train.cpu()
D_cost_valid = D_cost_valid.cpu()
D_wass_train = D_wass_train.cpu()
D_wass_valid = D_wass_valid.cpu()
batch_history['discriminator'].append({
'cost':
D_cost_train.data.numpy()[0],
'wasserstein_cost':
D_wass_train.data.numpy()[0],
'cost_validation':
D_cost_valid.data.numpy()[0],
'wasserstein_cost_validation':
D_wass_valid.data.numpy()[0]
})
############################
# (2) Update G network
###########################
# Prevent discriminator from computing gradients, since
# we are only updating the generator
for p in model_dis.parameters():
p.requires_grad = False
G_cost = compute_gener_loss_terms(model_dis,
model_gen,
batch_size,
latent_dim,
use_cuda,
compute_grads=True)
# Update generator
optimizer_gen.step()
if use_cuda:
G_cost = G_cost.cpu()
# Record generator loss
batch_history['generator']['cost'] = G_cost.data.numpy()[0]
# Record batch metrics
epoch_history.append(batch_history)
# Record epoch metrics
history.append(epoch_history)
LOGGER.info(pprint.pformat(epoch_history[-1]))
if (epoch + 1) % epochs_per_sample == 0:
# Generate outputs for fixed latent samples
LOGGER.info('Generating samples...')
samp_output = model_gen(sample_noise_v)
if use_cuda:
samp_output = samp_output.cpu()
samples[epoch + 1] = samp_output.data.numpy()
if output_dir:
LOGGER.info('Saving samples...')
save_samples(samples[epoch + 1], epoch + 1, output_dir)
## Get final discriminator loss
# Get noise
noise = torch.Tensor(batch_size, latent_dim).uniform_(-1, 1)
if use_cuda:
noise = noise.cuda()
noise_v = autograd.Variable(noise,
volatile=True) # totally freeze model_gen
# Get new batch of real training data
D_cost_test, D_wass_test = compute_discr_loss_terms(model_dis,
model_gen,
test_data_v,
noise_v,
batch_size,
latent_dim,
lmbda,
use_cuda,
compute_grads=False)
D_cost_valid, D_wass_valid = compute_discr_loss_terms(model_dis,
model_gen,
valid_data_v,
noise_v,
batch_size,
latent_dim,
lmbda,
use_cuda,
compute_grads=False)
if use_cuda:
D_cost_test = D_cost_test.cpu()
D_cost_valid = D_cost_valid.cpu()
D_wass_test = D_wass_test.cpu()
D_wass_valid = D_wass_valid.cpu()
final_discr_metrics = {
'cost_validation': D_cost_valid.data.numpy()[0],
'wasserstein_cost_validation': D_wass_valid.data.numpy()[0],
'cost_test': D_cost_test.data.numpy()[0],
'wasserstein_cost_test': D_wass_test.data.numpy()[0],
}
return model_gen, model_dis, history, final_discr_metrics, samples
def validate_decoder(target_path, generation_path):
from numpy import linalg as LA
import numpy as np
target_info, targets, _ = load_samples(target_path)
generate_info, generates, _ = load_samples(generation_path)
if len(targets) != len(generates):
return False
for i in range(len(targets)):
if LA.norm(np.subtract(targets[i], generates[i])) > 0.001:
return False
return True
if __name__ == "__main__":
test_info, test_data = load_latent_vectors(options.data_path)
model = load_model(options.model_path)
decoder = Model(inputs=model.get_layer('sequential_2').model.input,
outputs=model.get_layer('sequential_2').model.output)
generations = decoder.predict(test_data)
save_samples(os.path.join(options.save_path, 'generates'), generations,
test_info)
print(
validate_decoder(os.path.join(options.save_path, 'recons'),
os.path.join(options.save_path, 'generates')))
def train(self, dataset):
try:
temp = set(tf.global_variables())
except:
temp = set(tf.all_variables())
print("Creating optimizer")
d_optim = tf.train.AdamOptimizer(Options.lrate_d, beta1=Options.beta1_d) \
.minimize(self.d_loss, var_list=self.d_vars)
g_optim = tf.train.AdamOptimizer(Options.lrate_g, beta1=Options.beta1_g) \
.minimize(self.g_loss, var_list=self.g_vars)
print("Optimizer created")
with tf.Session() as sess:
self.sess = sess
try:
print("Restoring from checkpoint")
_dir = os.path.join(Options.checkpoint_dir, self.prefix())
a = glob.glob(os.path.join(_dir, '*/'))
a = [int(i.split('/')[-2]) for i in a]
a.sort()
_dir = os.path.join(_dir, str(a[-1]), 'model.ckpt')
self.saver.restore(self.sess, _dir)
print("Model restored")
try:
self.sess.run(
tf.variables_initializer(
set(tf.global_variables()) - temp))
except:
self.sess.run(
tf.initialize_variables(
set(tf.all_variables()) - temp))
except:
print("Initializing")
try:
tf.global_variables_initializer().run()
except:
tf.initialize_all_variables().run()
print("Initialized")
print("Merging summaries")
self.g_sum = md.merge_summary([
self.z_sum, self.d__sum, self.d_loss_fake_sum, self.g_loss_sum
])
self.d_sum = md.merge_summary([
self.z_sum, self.d_sum, self.d_loss_real_sum, self.d_loss_sum
])
self.writer = md.SummaryWriter("./logs", self.sess.graph)
print("Summaries merged")
sample_z = np.random.uniform(-1,
1,
size=(self.sample_size, self.z_dim))
counter = 0
terrD_fake = 0.0
terrD_real = 0.0
terrG = 0.0
s_begin = time.time()
c_begin = time.time()
p_begin = time.time()
print("Starting training epoch")
for epoch in range(Options.train_epochs):
for sub_data in dataset.train_iter():
batch_z = np.random.uniform(
-1, 1,
[Options.batch_size, self.z_dim]).astype(np.float32)
# Update D network
_, summary_str = self.sess.run([d_optim, self.d_sum],
feed_dict={
self.videos: sub_data,
self.z: batch_z
})
self.writer.add_summary(summary_str, counter)
# Update G network
_, summary_str = self.sess.run([g_optim, self.g_sum],
feed_dict={self.z: batch_z})
self.writer.add_summary(summary_str, counter)
errD_fake = self.d_loss_fake.eval({self.z: batch_z})
errD_real = self.d_loss_real.eval({self.videos: sub_data})
errG = self.g_loss.eval({self.z: batch_z})
terrD_real += errD_real
terrD_fake += errD_fake
terrG += errG
counter += 1
if time.time() - p_begin > Options.print_time:
print(
"Epoch: [%d], d_loss_fake: [%.6f]--[%.4f], d_loss_real: [%.6f]--[%.4f], g_loss: [%.6f]--[%.4f]"
% (epoch, terrD_fake / counter, errD_fake,
terrD_real / counter, errD_real,
terrG / counter, errG))
p_begin = time.time()
if time.time() - s_begin > Options.sampler_time:
samples, d_loss, g_loss = self.sess.run(
[self.sampler, self.d_loss, self.g_loss],
feed_dict={
self.z: sample_z,
self.videos: sub_data
})
utils.save_samples(samples, epoch, counter)
print("[Sample] d_loss: %.8f, g_loss: %.8f" %
(d_loss, g_loss))
s_begin = time.time()
if time.time() - c_begin > Options.checkpoint_time:
print("Checkpointing")
if not os.path.exists(Options.checkpoint_dir):
os.makedirs(Options.checkpoint_dir)
_dir = os.path.join(Options.checkpoint_dir,
self.prefix())
self.save(_dir, epoch, self.sess)
c_begin = time.time()
counter = 0
terrD_fake = 0.0
terrD_real = 0.0
terrG = 0.0
def train(self):
"""
Draws samples from the data distribution and the noise distribution,
and alternates between optimizing the parameters of the discriminator
and the generator.
"""
# TODO add loop over epochs and batches instead single iterations
with tf.Session() as session:
tf.global_variables_initializer().run()
if pre_train:
# pretraining discriminator
for step in range(self.pre_train_steps):
# d = (np.random.random(self.batch_size) - 0.5) * 10.0
d = self.data[self.pre_train_steps + step]
# labels = norm.pdf(d, loc=self.data.mu, scale=self.data.sigma)
# labels = self.data[step]
pretrain_loss, _ = session.run(
[self.pre_loss, self.pre_opt], {
self.pre_input: np.reshape(d, self.input_shape)
})
self.weightsD = session.run(self.d_pre_params)
# copy weights from pre-training over to new D network
for i, v in enumerate(self.d_params):
session.run(tf.assign(v, self.weightsD[i]))
# session.run(v.assign(self.weightsD[i]))
writer = tf.summary.FileWriter("./logs", session.graph)
g_sum = tf.summary.merge(
[self.z_sum, self.d2_sum, self.G_sum, self.loss_d_fake_sum,
self.loss_g_sum])
d_sum = tf.summary.merge(
[self.z_sum, self.d1_sum, self.loss_d_real_sum,
self.loss_d_sum])
for step in range(self.num_steps):
print('iteration at step {}'.format(step))
# update discriminator
# x = np.sort(self.data[self.num_steps + step])
x = self.data[step]
# TODO try different noise sources
# z = self.gen.sample_int(self.input_shape[0])
z = self.gen.binned_samples(self.input_shape)
loss_d, summary, _ = session.run(
[self.loss_d, d_sum, self.opt_d],
{
self.x: x[np.newaxis, :, :, np.newaxis],
self.z: z,
self.is_training: True
})
writer.add_summary(summary, step)
# update generator
# z = self.gen.sample_int(self.input_shape[0])
z = self.gen.binned_samples(self.input_shape)
loss_g, summary, _ = session.run(
[self.loss_g, g_sum, self.opt_g],
{
self.z: np.reshape(z, self.input_shape),
self.is_training: True
})
writer.add_summary(summary, step)
errD_fake = self.loss_d_fake.eval(
{self.z: z, self.is_training: False})
errD_real = self.loss_d_real.eval(
{self.x: x[np.newaxis, :, :, np.newaxis], self.is_training: False})
errG = self.loss_g.eval(
{self.z: z, self.is_training: False})
print('loss_d {}, errD_real {}'.format(loss_d, errD_real + errD_fake))
print('loss_g {}, errG {}'.format(loss_g, errG))
self.loss_d_plot.append(loss_d)
self.loss_g_plot.append(loss_g)
if step % self.log_every == 0:
# print('{}: loss_d: {}\t loss_g: {}'.format(step, loss_d,
# loss_g))
samples, d_loss, g_loss = session.run(
[self.G, self.loss_d, self.loss_g],
feed_dict={self.z: z,
self.x: x[np.newaxis, :, :, np.newaxis],
self.is_training: False}
)
now = datetime.datetime.now().isoformat()
save_samples(samples, self.samples_dir,
'train_{}_{}.npy'.format(step, now))
if self.anim_path:
self.anim_frames.append(plots.samples(session, save=False,
lower_range=self.gen.lower_range,
upper_range=self.gen.upper_range,
batch_size=self.input_shape[0],
D1=self.D1,
G=self.G,
x=self.x,
data=self.data,
z=self.z))
if (step % 500) == 0:
save(self.checkpoint_dir, step, self.saver, session,
'GAN.model')
if self.anim_path:
plots.save_animation(self.anim_path, self.anim_frames,
lower_range=self.gen.lower_range,
upper_range=self.gen.upper_range)
else:
# TODO find a proper plotting mechanism
pass
# plots.plot_distributions(session, save=False,
# lower_range=self.gen.lower_range,
# upper_range=self.gen.upper_range,
# batch_size=self.input_shape[0],
# D1=self.D1,
# G=self.G,
# x=self.x,
# z=self.z,
# data=self.data)
# plots.plot_training_loss(self.loss_g_plot, self.loss_d_plot)
# merged = tf.summary.merge_all()
# writer = tf.summary.FileWriter('logs', session.graph)
writer.close()
def main():
data_name = "augmented_1"
data_path = os.path.join("./data", data_name)
csv_name = data_name + ".csv"
train_df = pd.read_csv(os.path.join(data_path, csv_name))
keypoint_names = list(
map(lambda x: x[:-2],
train_df.columns.to_list()[1::2]))
keypoint_flip_map = [
("left_eye", "right_eye"),
("left_ear", "right_ear"),
("left_shoulder", "right_shoulder"),
("left_elbow", "right_elbow"),
("left_wrist", "right_wrist"),
("left_hip", "right_hip"),
("left_knee", "right_knee"),
("left_ankle", "right_ankle"),
("left_palm", "right_palm"),
("left_instep", "right_instep"),
]
image_list = train_df.iloc[:, 0].to_numpy()
keypoints_list = train_df.iloc[:, 1:].to_numpy()
train_imgs, valid_imgs, train_keypoints, valid_keypoints = train_val_split(
image_list, keypoints_list, random_state=42)
image_set = {"train": train_imgs, "valid": valid_imgs}
keypoints_set = {"train": train_keypoints, "valid": valid_keypoints}
hyper_params = {
"augmented_ver": data_name,
"learning_rate": 0.001,
"num_epochs": 10000,
"batch_size": 256,
"description": "Final training"
}
for phase in ["train", "valid"]:
DatasetCatalog.register(
"keypoints_" + phase,
lambda phase=phase: get_data_dicts(data_path, image_set[phase],
keypoints_set[phase]))
MetadataCatalog.get("keypoints_" + phase).set(thing_classes=["human"])
MetadataCatalog.get("keypoints_" +
phase).set(keypoint_names=keypoint_names)
MetadataCatalog.get("keypoints_" +
phase).set(keypoint_flip_map=keypoint_flip_map)
MetadataCatalog.get("keypoints_" + phase).set(evaluator_type="coco")
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-Keypoints/keypoint_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("keypoints_train", )
cfg.DATASETS.TEST = ("keypoints_valid", )
cfg.DATALOADER.NUM_WORKERS = 16 # On Windows environment, this value must be 0.
cfg.SOLVER.IMS_PER_BATCH = 2 # mini batch size would be (SOLVER.IMS_PER_BATCH) * (ROI_HEADS.BATCH_SIZE_PER_IMAGE).
cfg.SOLVER.BASE_LR = hyper_params["learning_rate"] # Learning Rate.
cfg.SOLVER.MAX_ITER = hyper_params["num_epochs"] # Max iteration.
cfg.SOLVER.GAMMA = 0.8
cfg.SOLVER.STEPS = [
3000, 4000, 5000, 6000, 7000, 8000
] # The iteration number to decrease learning rate by GAMMA.
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-Keypoints/keypoint_rcnn_X_101_32x8d_FPN_3x.yaml")
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = hyper_params[
"batch_size"] # Use to calculate RPN loss.
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS = 24
cfg.TEST.KEYPOINT_OKS_SIGMAS = np.ones((24, 1), dtype=float).tolist()
cfg.TEST.EVAL_PERIOD = 5000 # Evaluation would occur for every cfg.TEST.EVAL_PERIOD value.
cfg.OUTPUT_DIR = os.path.join("./output", data_name)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = Trainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
# Inference should use the config with parameters that are used in training
# cfg now already contains everything we've set previously. We changed it a little bit for inference:
cfg.MODEL.WEIGHTS = os.path.join(
cfg.OUTPUT_DIR, "model_final.pth") # path to the model we just trained
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set a custom testing threshold
predictor = DefaultPredictor(cfg)
test_dir = os.path.join("data", "test_imgs")
test_list = os.listdir(test_dir)
test_list.sort()
except_list = []
files = []
preds = []
for file in tqdm(test_list):
filepath = os.path.join(test_dir, file)
im = cv2.imread(filepath)
outputs = predictor(im)
outputs = outputs["instances"].to("cpu").get("pred_keypoints").numpy()
files.append(file)
pred = []
try:
for out in outputs[0]:
pred.extend([float(e) for e in out[:2]])
except IndexError:
pred.extend([0] * 48)
except_list.append(filepath)
preds.append(pred)
df_sub = pd.read_csv("./data/sample_submission.csv")
df = pd.DataFrame(columns=df_sub.columns)
df["image"] = files
df.iloc[:, 1:] = preds
df.to_csv(os.path.join(cfg.OUTPUT_DIR, f"{data_name}_submission.csv"),
index=False)
if except_list:
print(
"The following images are not detected keypoints. The row corresponding that images names would be filled with 0 value."
)
print(*except_list)
save_samples(cfg.OUTPUT_DIR,
test_dir,
os.path.join(cfg.OUTPUT_DIR, f"{data_name}_submission.csv"),
mode="random",
size=5)
def train(opt, dataloader_X, dataloader_Y, test_dataloader_X,
test_dataloader_Y):
input_shape = (None, None, opt["channels"])
norm_type = 'batchnorm'
if opt["batch_size"] == 1:
norm_type = 'instancenorm'
# models
G_YtoX = model.Generator(input_shape, opt, norm_type)
G_XtoY = model.Generator(input_shape, opt, norm_type)
D_X = model.Discriminator(input_shape, norm_type)
D_Y = model.Discriminator(input_shape, norm_type)
# summary
G_YtoX.summary()
D_X.summary()
# optimizers
G_YtoX_optimizer = tf.optimizers.Adam(opt["lr_g"], beta_1=opt["b1"])
G_XtoY_optimizer = tf.optimizers.Adam(opt["lr_g"], beta_1=opt["b1"])
D_X_optimizer = tf.optimizers.Adam(opt["lr_d"], beta_1=opt["b1"])
D_Y_optimizer = tf.optimizers.Adam(opt["lr_d"], beta_1=opt["b1"])
# set checkpoint
ckpt, ckpt_manager = model.set_checkpoint(opt, G_YtoX, G_XtoY, D_X, D_Y,
G_YtoX_optimizer,
G_XtoY_optimizer, D_X_optimizer,
D_Y_optimizer)
test_X = next(iter(test_dataloader_X))
test_Y = next(iter(test_dataloader_Y))
iterations = 0
total_disc_loss_list, total_gen_loss_list, test_loss_list = utils.load_losses_list(
opt)
mean_cycle_loss_test = 0
count = 0
for epoch in range(opt["epoch"], opt["n_epochs"] + 1):
start = time.time()
for image_x, image_y in tf.data.Dataset.zip(
(dataloader_X, dataloader_Y)):
total_disc_loss, total_gen_loss = train_step(
image_x, image_y, G_YtoX, G_XtoY, D_X, D_Y, G_YtoX_optimizer,
G_XtoY_optimizer, D_X_optimizer, D_Y_optimizer, opt)
total_disc_loss_list.append(total_disc_loss)
total_gen_loss_list.append(total_gen_loss)
if iterations % 200 == 0:
print('iteration:{}'.format(iterations))
# Using a consistent image so that the progress of the model is visible.
cycle_loss_test = utils.save_samples(epoch, test_Y, test_X,
G_YtoX, G_XtoY, opt)
mean_cycle_loss_test += cycle_loss_test
count += 1
utils.plot_losses(total_disc_loss_list, total_gen_loss_list,
opt, epoch, "Epoch:{}".format(epoch))
iterations += 1
clear_output(wait=True)
# Using a consistent image so that the progress of the model is visible. Only one plot is keept per epoch
cycle_loss_test = utils.save_samples(epoch, test_Y, test_X, G_YtoX,
G_XtoY, opt)
mean_cycle_loss_test += cycle_loss_test
count += 1
test_loss_list.append(mean_cycle_loss_test / count)
mean_cycle_loss_test = 0
count = 0
utils.plot_losses(total_disc_loss_list, total_gen_loss_list, opt,
epoch, "Epoch:{}".format(epoch))
utils.plot_test_loss(test_loss_list, opt, epoch)
utils.save_losses_list(total_disc_loss_list, total_gen_loss_list,
test_loss_list, opt)
print("Sample saved...")
if (epoch + 1) % opt["checkpoint_interval"] == 0:
ckpt_save_path = ckpt_manager.save()
print('Saving checkpoint for epoch {} at {}'.format(
epoch, ckpt_save_path))
print('Time taken for epoch {} is {} sec\n'.format(
epoch,
time.time() - start))
return G_YtoX, G_XtoY, D_X, D_Y
norm_data = np.empty((len(encoded_data), 1, imageSize, imageSize),
dtype=np.float32)
norm_value = []
for i, ed in enumerate(encoded_data):
dat = np.array(ed).reshape((1, imageSize, imageSize))
norm_data[i] = np.divide(dat, np.max(dat))
norm_value.append(np.max(dat))
save_dict['normed_data'] = norm_data
save_dict['num_samples'] = num_samples
save_dict['imageSize'] = imageSize
save_dict['spikes'] = raw_data[:100]
save_dict['encoded_data'] = encoded_data
save_dict['normed_values'] = norm_value
else:
save_dict['binned_data'] = binned_data
save_dict['normed_data'] = norm_data
save_dict['num_samples'] = num_samples
save_dict['imageSize'] = imageSize
save_dict['spikes'] = raw_data
save_dict['data_type'] = opt.data_type
if opt.filename:
fname = opt.filename
else:
fname = 'data_NS{}_IS{}_type-{}_encoded-{}_rate{}.npy'.format(
num_samples, imageSize, opt.data_type, opt.encoding, ARRAY_ID)
u = str(uuid.uuid4())
path = os.path.join(opt.path, u)
utils.save_samples(save_dict, path=path, filename=fname)
def train(self):
opti_D = tf.train.AdamOptimizer(learning_rate=self.learn_rate,
beta1=0.5).minimize(
self.D_loss, var_list=self.d_var)
opti_G = tf.train.AdamOptimizer(learning_rate=self.learn_rate,
beta1=0.5).minimize(
self.G_loss, var_list=self.g_var)
init = tf.global_variables_initializer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(init)
summary_writer = tf.summary.FileWriter(self.log_dir,
graph=sess.graph)
step = 0
while step <= 10000:
realbatch_array, real_labels = self.data_ob.getNext_batch(step)
# Get the z
batch_z = np.random.uniform(-1,
1,
size=[self.batch_size, self.z_dim])
_, summary_str = sess.run(
[opti_D, self.merged_summary_op_d],
feed_dict={
self.real_feature_vector: realbatch_array,
self.z: batch_z,
self.y: real_labels
})
summary_writer.add_summary(summary_str, step)
_, summary_str = sess.run([opti_G, self.merged_summary_op_g],
feed_dict={
self.z: batch_z,
self.y: real_labels
})
summary_writer.add_summary(summary_str, step)
if step % 50 == 0:
D_loss = sess.run(self.D_loss,
feed_dict={
self.real_feature_vector:
realbatch_array,
self.z: batch_z,
self.y: real_labels
})
fake_loss = sess.run(self.G_loss,
feed_dict={
self.z: batch_z,
self.y: real_labels
})
print("Step %d: D: loss = %.7f G: loss=%.7f " %
(step, D_loss, fake_loss))
if np.mod(step, 50) == 1 and step != 0:
sample_exec_time = sample_output(
batch_size=self.batch_size,
min=self.y_min,
max=self.y_max)
sample_feature = sess.run(self.fake_feature_vector,
feed_dict={
self.z: batch_z,
self.y: sample_exec_time
})
save_samples(sample_feature, sample_exec_time)
self.saver.save(sess, self.model_path)
step = step + 1
save_path = self.saver.save(sess, self.model_path)
print("Model saved in file: %s" % save_path)
def train(self, ImageModel, netG, netD, audio_encoder, ImageGeneratorModel,
ImageDiscriminatorModel, args):
if self.use_cuda:
netG.cuda()
netD.cuda()
ImageModel.cuda()
audio_encoder.cuda()
ImageGeneratorModel.cuda()
ImageDiscriminatorModel.cuda()
logger = Logger(self.log_folder)
##audio part
optimizerG = optim.Adam(netG.parameters(),
lr=float(args['--learning_rate']),
betas=(float(args['--beta1']),
float(args['--beta2'])))
optimizerD = optim.Adam(netD.parameters(),
lr=float(args['--learning_rate']),
betas=(float(args['--beta1']),
float(args['--beta2'])))
opt_generator = optim.Adam(ImageModel.parameters(),
lr=0.0002,
betas=(0.5, 0.999),
weight_decay=0.00001)
optimizerG_Image = optim.Adam(ImageGeneratorModel.parameters(),
lr=float(args['--learning_rate']),
betas=(float(args['--beta1']),
float(args['--beta2'])))
optimizerD_image = optim.Adam(ImageDiscriminatorModel.parameters(),
lr=float(args['--learning_rate']),
betas=(float(args['--beta1']),
float(args['--beta2'])))
opt_encoder = optim.Adam(audio_encoder.parameters(),
lr=float(args['--learning_rate']),
betas=(float(args['--beta1']),
float(args['--beta2'])))
# opt_encoder = optim.Adam(audio_encoder.parameters(), lr=0.0002, betas=(0.5, 0.999), weight_decay=0.00001)
# =============Train===============
lmbda = float(args['--lmbda'])
epochs_per_sample = int(args['--epochs_per_sample'])
output_dir = 'outputs'
history = []
D_costs_train = []
D_wasses_train = []
D_costs_valid = []
D_wasses_valid = []
G_costs = []
BATCH_NUM = 350
epochs = 1800
start = time.time()
self.LOGGER_audio.info(
'Starting training...EPOCHS={}, BATCH_SIZE={}, BATCH_NUM={}'.
format(epochs, self.image_batch_size, BATCH_NUM))
self.LOGGER_image.info(
'Starting training...EPOCHS={}, BATCH_SIZE={}, BATCH_NUM={}'.
format(epochs, self.audio_batch_size, BATCH_NUM))
# training loop
def sample_fake_image_batch(batch_size, real_img):
return generator.sample_images(batch_size, real_img)
def sample_fake_video_batch(batch_size, real_video):
return generator.sample_videos(batch_size, real_video)
def init_logs():
return {'l_gen': 0, 'l_image_dis': 0, 'l_video_dis': 0}
batch_num = 0
logs = init_logs()
start_time = time.time()
epoch = 0
gc.collect()
while True:
gc.collect()
epoch = epoch + 1
self.LOGGER_audio.info("{} Epoch: {}/{}".format(
time_since(start), epoch, epochs))
D_cost_train_epoch = []
D_wass_train_epoch = []
D_cost_valid_epoch = []
D_wass_valid_epoch = []
G_cost_epoch = []
self.LOGGER_image.info("{} Epoch: {}/{}".format(
time_since(start), epoch, epochs))
D2_cost_train_epoch = []
D2_wass_train_epoch = []
D2_cost_valid_epoch = []
D2_wass_valid_epoch = []
G2_cost_epoch = []
for i in range(1, BATCH_NUM + 1):
# sample real data
sample_real_batch = self.sample_real_image_batch()
batch_real_audio = sample_real_batch['audio'].cpu()
batch_image = Variable(sample_real_batch['images'],
requires_grad=False)
# print(batch_image.shape)
batch_image = batch_image.cuda()
for p in netD.parameters():
p.requires_grad = True
one = torch.Tensor([1]).float()
neg_one = one * -1
if self.use_cuda:
one = one.cuda()
neg_one = neg_one.cuda()
# (1) Train Discriminator
for iter_dis in range(1):
netD.zero_grad()
z = nn.init.normal(torch.Tensor(self.image_batch_size,
512))
if self.use_cuda:
z = z.cuda()
z = Variable(z)
real_data_Var = numpy_to_var(batch_real_audio,
self.use_cuda)
# print(batch_image.shape)
# print(type(batch_image))
batch_image = batch_image.cuda()
real_img_Var = Variable(batch_image)
# a) compute loss contribution from real training data
D_real = netD(real_data_Var)
D_real = D_real.mean()
#print('D_real',D_real) # avg loss
D_real.backward(neg_one) # loss * -1
#print('real_img_Var',real_img_Var.shape)
D_real_img = ImageDiscriminatorModel(real_img_Var)
D_real_img = D_real_img.mean() # avg loss
#print('D_real_img',D_real_img)
D_real_img = Variable(D_real_img.data,
requires_grad=True) #Added
D_real_img.backward(neg_one) # loss * -1
# b) compute loss contribution from generated data, then backprop.
features = ImageModel(batch_image)
fk_audio = netG(z, features)
fk_audio = autograd.Variable(fk_audio.data)
#print(fk_audio.shape)(16, 1, 16384)
D_fake = netD(fk_audio)
D_fake = D_fake.mean()
D_fake.backward(one)
#print(batch_real_audio.shape)#16*16384
batch_real_audio = batch_real_audio.unsqueeze(1)
#print(batch_real_audio.shape)#16*1*16384
# print(fk_audio.shape)
# print(type(fk_audio))
# print(batch_real_audio.shape)
# print(type(batch_real_audio))
# print('audio_encoder',audio_encoder)
audio_features = audio_encoder(batch_real_audio)
#print('audio_features',audio_features.shape)
#audio_features=audio_encoder(fk_audio)
audio_features = audio_features.unsqueeze(2).unsqueeze(3)
fk_image = ImageGeneratorModel(audio_features)
fk_image = autograd.Variable(fk_image.data)
D_fake_image = ImageDiscriminatorModel(fk_image)
D_fake_image = D_fake_image.mean()
D_fake_image = Variable(D_fake_image.data,
requires_grad=True) #Added
D_fake_image.backward(one)
#print('real_img_Var',type(real_img_Var.data))#16x3x224x224
# c) compute gradient penalty and backprop
gradient_penalty = calc_gradient_penalty(
netD,
real_data_Var.data,
fk_audio.data,
self.image_batch_size,
lmbda,
use_cuda=self.use_cuda)
gradient_penalty.backward(one)
################################# 16*3*224*224 to 16*3*64*64 or do nn.AvgPool2d
LRTrans = transforms.Compose([
transforms.Scale(64, Image.BICUBIC),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
real_img_Var_64 = torch.zeros([16, 3, 64, 64])
for j in range(16):
#print('j',j)
real_img_Var_64[j] = torch.FloatTensor(
LRTrans(
Image.fromarray(real_img_Var.data.cpu().mul(
0.5).add(0.5).mul(255).byte()[j].transpose(
0, 2).transpose(0,
1).numpy())).numpy())
#print('testing',real_img_Var.data.cpu().mul(0.5).add(0.5).mul(255).byte()[j].shape)
#print("real_img_Var_64",real_img_Var_64.shape)
real_img_Var_64 = real_img_Var_64.cuda()
######################################
gradient_penalty_2 = calc_gradient_penalty_2(
ImageDiscriminatorModel,
real_img_Var_64.data,
fk_image.data,
self.audio_batch_size,
lmbda,
use_cuda=self.use_cuda)
gradient_penalty_2.backward(one)
# Compute cost * Wassertein loss..
D_cost_train = D_fake - D_real + gradient_penalty
D_wass_train = D_real - D_fake
D2_cost_train = D_fake_image - D_real_img + gradient_penalty_2
D2_wass_train = D_real_img - D_fake_image
# Update gradient of discriminator.
optimizerD.step()
optimizerD_image.step()
#############################
# (2) Compute Valid data
#############################
netD.zero_grad()
ImageDiscriminatorModel.zero_grad()
batch_real_audio = batch_real_audio.squeeze(1)
valid_data_Var = numpy_to_var(batch_real_audio,
self.use_cuda)
D_real_valid = netD(valid_data_Var)
D_real_valid = D_real_valid.mean() # avg loss
#valid_data_Var_2 = numpy_to_var(batch_image, self.use_cuda)# can substitute this with below two lines
batch_image = batch_image.cuda()
valid_data_Var_2 = Variable(batch_image)
D2_real_valid = ImageDiscriminatorModel(valid_data_Var_2)
D2_real_valid = D2_real_valid.mean() # avg loss
# b) compute loss contribution from generated data, then backprop.
fake_valid = netG(z, features)
D_fake_valid = netD(fake_valid)
D_fake_valid = D_fake_valid.mean()
fake_valid_2 = ImageGeneratorModel(audio_features)
D2_fake_valid = ImageDiscriminatorModel(fake_valid_2)
D2_fake_valid = D2_fake_valid.mean()
# c) compute gradient penalty and backprop
gradient_penalty_valid = calc_gradient_penalty(
netD,
valid_data_Var.data,
fake_valid.data,
self.image_batch_size,
lmbda,
use_cuda=self.use_cuda)
################################# 16*3*224*224 to 16*3*64*64 or do nn.AvgPool2d
LRTrans = transforms.Compose([
transforms.Scale(64, Image.BICUBIC),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
valid_data_Var_2_64 = torch.zeros([16, 3, 64, 64])
for j in range(16):
valid_data_Var_2_64[j] = torch.FloatTensor(
LRTrans(
Image.fromarray(
valid_data_Var_2.data.cpu().mul(0.5).add(
0.5).mul(255).byte()[j].transpose(
0,
2).transpose(0,
1).numpy())).numpy())
#print('testing',real_img_Var.data.cpu().mul(0.5).add(0.5).mul(255).byte()[j].shape)
#print("valid_data_Var_2_64",valid_data_Var_2_64.shape)
valid_data_Var_2_64 = valid_data_Var_2_64.cuda()
######################################
gradient_penalty_valid_2 = calc_gradient_penalty_2(
ImageDiscriminatorModel,
valid_data_Var_2_64.data,
fake_valid_2.data,
self.audio_batch_size,
lmbda,
use_cuda=self.use_cuda)
# Compute metrics and record in batch history.
D_cost_valid = D_fake_valid - D_real_valid + gradient_penalty_valid
D_wass_valid = D_real_valid - D_fake_valid
D2_cost_valid = D2_fake_valid - D2_real_valid + gradient_penalty_valid_2
D2_wass_valid = D2_real_valid - D2_fake_valid
if self.use_cuda:
D_cost_train = D_cost_train.cpu()
D_wass_train = D_wass_train.cpu()
D_cost_valid = D_cost_valid.cpu()
D_wass_valid = D_wass_valid.cpu()
D2_cost_train = D2_cost_train.cpu()
D2_wass_train = D2_wass_train.cpu()
D2_cost_valid = D2_cost_valid.cpu()
D2_wass_valid = D2_wass_valid.cpu()
# Record costs
D_cost_train_epoch.append(D_cost_train.data.numpy())
D_wass_train_epoch.append(D_wass_train.data.numpy())
D_cost_valid_epoch.append(D_cost_valid.data.numpy())
D_wass_valid_epoch.append(D_wass_valid.data.numpy())
D2_cost_train_epoch.append(D2_cost_train.data.numpy())
D2_wass_train_epoch.append(D2_wass_train.data.numpy())
D2_cost_valid_epoch.append(D2_cost_valid.data.numpy())
D2_wass_valid_epoch.append(D2_wass_valid.data.numpy())
#############################
# (3) Train Generator
#############################
# Prevent discriminator update.
for p in netD.parameters():
p.requires_grad = False
for p in ImageDiscriminatorModel.parameters():
p.requires_grad = False
# Reset generator gradients
netG.zero_grad()
fk_audio = netG(z, features)
# fake = autograd.Variable(fk_img.data)
# fake = netG(fk_img2)
G = netD(fk_audio)
G = G.mean()
# print('audio_cond',audio_cond.shape)
# Update gradients.
G.backward(neg_one)
G_cost = -G
optimizerG.step()
opt_generator.step()
ImageGeneratorModel.zero_grad()
fk_img = ImageGeneratorModel(audio_features)
# fake = autograd.Variable(fk_img.data)
# fake = netG(fk_img2)
G2 = ImageDiscriminatorModel(fk_img)
G2 = G2.mean()
# print('audio_cond',audio_cond.shape)
# Update gradients.
G2.backward(neg_one)
G2_cost = -G2
optimizerG_Image.step()
opt_encoder.step()
# Record costs
if self.use_cuda:
G_cost = G_cost.cpu()
G_cost_epoch.append(G_cost.data.numpy())
if i % (BATCH_NUM // 5) == 0:
self.LOGGER_audio.info(
"{} Epoch={} Batch: {}/{} D_c:{:.4f} | D_w:{:.4f} | G:{:.4f}"
.format(time_since(start), epoch, i, BATCH_NUM,
D_cost_train.data.numpy(),
D_wass_train.data.numpy(),
G_cost.data.numpy()))
if self.use_cuda:
G2_cost = G2_cost.cpu()
G2_cost_epoch.append(G2_cost.data.numpy())
if i % (BATCH_NUM // 5) == 0:
self.LOGGER_image.info(
"{} Epoch={} Batch: {}/{} D2_c:{:.4f} | D2_w:{:.4f} | G2:{:.4f}"
.format(time_since(start), epoch, i, BATCH_NUM,
D2_cost_train.data.numpy(),
D2_wass_train.data.numpy(),
G2_cost.data.numpy()))
# Save the average cost of batches in every epoch.
D_cost_train_epoch_avg = sum(D_cost_train_epoch) / float(
len(D_cost_train_epoch))
D_wass_train_epoch_avg = sum(D_wass_train_epoch) / float(
len(D_wass_train_epoch))
D_cost_valid_epoch_avg = sum(D_cost_valid_epoch) / float(
len(D_cost_valid_epoch))
D_wass_valid_epoch_avg = sum(D_wass_valid_epoch) / float(
len(D_wass_valid_epoch))
G_cost_epoch_avg = sum(G_cost_epoch) / float(len(G_cost_epoch))
D_costs_train.append(D_cost_train_epoch_avg)
D_wasses_train.append(D_wass_train_epoch_avg)
D_costs_valid.append(D_cost_valid_epoch_avg)
D_wasses_valid.append(D_wass_valid_epoch_avg)
G_costs.append(G_cost_epoch_avg)
# Save the average cost of batches in every epoch.
D2_cost_train_epoch_avg = sum(D2_cost_train_epoch) / float(
len(D2_cost_train_epoch))
D2_wass_train_epoch_avg = sum(D2_wass_train_epoch) / float(
len(D2_wass_train_epoch))
D2_cost_valid_epoch_avg = sum(D2_cost_valid_epoch) / float(
len(D2_cost_valid_epoch))
D2_wass_valid_epoch_avg = sum(D2_wass_valid_epoch) / float(
len(D2_wass_valid_epoch))
G2_cost_epoch_avg = sum(G2_cost_epoch) / float(len(G2_cost_epoch))
D2_costs_train.append(D2_cost_train_epoch_avg)
D2_wasses_train.append(D2_wass_train_epoch_avg)
D2_costs_valid.append(D2_cost_valid_epoch_avg)
D2_wasses_valid.append(D2_wass_valid_epoch_avg)
G2_costs.append(G2_cost_epoch_avg)
self.LOGGER_audio.info(
"{} D_cost_train:{:.4f} | D_wass_train:{:.4f} | D_cost_valid:{:.4f} | D_wass_valid:{:.4f} | "
"G_cost:{:.4f}".format(time_since(start),
D_cost_train_epoch_avg,
D_wass_train_epoch_avg,
D_cost_valid_epoch_avg,
D_wass_valid_epoch_avg,
G_cost_epoch_avg))
self.LOGGER_image.info(
"{} D2_cost_train:{:.4f} | D2_wass_train:{:.4f} | D2_cost_valid:{:.4f} | D2_wass_valid:{:.4f} | "
"G2_cost:{:.4f}".format(time_since(start),
D2_cost_train_epoch_avg,
D2_wass_train_epoch_avg,
D2_cost_valid_epoch_avg,
D2_wass_valid_epoch_avg,
G2_cost_epoch_avg))
# Generate audio samples.
if epoch % epochs_per_sample == 0:
self.LOGGER_audio.info("Generating samples...")
self.LOGGER_image.info("Generating image samples...")
# batch_real_image_val,image_enumerator=sample_real_image_batch(image_enumerator)
# batch_real_audio_val=batch_real_image_val['audio'].cpu()
# sample_test_Var = numpy_to_var(batch_real_audio_val, cuda)
torch.save(
ImageModel,
os.path.join(self.log_folder,
'%05d_ImageModel.pytorch' % epoch))
torch.save(
netG,
os.path.join(self.log_folder, '%05d_netG.pytorch' % epoch))
torch.save(
netD,
os.path.join(self.log_folder,
'I%05d_netD.pytorch' % epoch))
torch.save(
audio_encoder,
os.path.join(self.log_folder,
'%05d_audio_encoder.pytorch' % epoch))
torch.save(
ImageGeneratorModel,
os.path.join(self.log_folder,
'%05d_ImageGeneratorModel.pytorch' % epoch))
torch.save(
ImageDiscriminatorModel,
os.path.join(
self.log_folder,
'I%05d_ImageDiscriminatorModel.pytorch' % epoch))
for iii in range(1, 2):
# sample real data
sample_test_batch = self.sample_test_image_batch()
batch_test_audio = sample_test_batch['audio'].cpu()
batch_image_test = Variable(sample_test_batch['images'],
requires_grad=False)
# print(batch_image.shape)
batch_image_test = batch_image_test.cuda()
features_test = ImageModel(batch_image)
z = nn.init.normal(torch.Tensor(self.image_batch_size,
512))
if self.use_cuda:
z = z.cuda()
z = Variable(z)
sample_out = netG(z, features_test)
if self.use_cuda:
sample_out = sample_out.cpu()
sample_out = sample_out.data.numpy()
save_samples(sample_out, epoch, output_dir)
gc.collect()
