batch_size=arglist.batch_size,
learning_rate=arglist.lr,
kl_tolerance=arglist.kl_tolerance,
is_training=True,
reuse=False,
gpu_mode=True)
# train loop:
print("train", "step", "loss", "recon_loss", "kl_loss")
for epoch in range(arglist.epoch):
np.random.shuffle(dataset)
for idx in range(num_batches):
batch = dataset[idx*arglist.batch_size:(idx+1)*arglist.batch_size]
obs = np.array(batch).astype(np.float)/255.0
feed = {vae.x: obs,}
(train_loss, r_loss, kl_loss, train_step, _) = vae.sess.run([
vae.loss, vae.r_loss, vae.kl_loss, vae.global_step, vae.train_op
], feed)
print("step", (train_step+1), train_loss, r_loss, kl_loss)
log.write('{}\n'.format(','.join(map(str,
[train_loss, r_loss, kl_loss, train_step]))))
if ((train_step+1) % arglist.save_period == 0):
vae.save_json(os.path.join(arglist.model_save_path,"vae.json"))
# finished, final model:
vae.save_json(os.path.join(arglist.model_save_path,"vae.json"))
vae = ConvVAE(z_size=z_size,
batch_size=batch_size,
learning_rate=learning_rate,
kl_tolerance=kl_tolerance,
is_training=True,
reuse=False,
gpu_mode=True)
# train loop:
print("train", "step", "loss", "recon_loss", "kl_loss")
for epoch in range(NUM_EPOCH):
np.random.shuffle(dataset)
for idx in range(num_batches):
batch = dataset[idx*batch_size:(idx+1)*batch_size]
obs = batch.astype(np.float)
feed = {vae.x: obs,}
(train_loss, r_loss, kl_loss, train_step, _) = vae.sess.run([
vae.loss, vae.r_loss, vae.kl_loss, vae.global_step, vae.train_op
], feed)
if ((train_step+1) % 500 == 0):
print("step", (train_step+1), train_loss, r_loss, kl_loss)
if ((train_step+1) % 5000 == 0):
vae.save_json("tf_vae/vae.json")
# finished, final model:
vae.save_json("tf_vae/vae.json")
def main( dirs, z_size=32, batch_size=100, learning_rate=0.0001, kl_tolerance=0.5, epochs=100, save_model=False, verbose=True, optimizer="Adam" ):
if save_model:
model_save_path = "tf_vae"
if not os.path.exists(model_save_path):
os.makedirs(model_save_path)
gen = DriveDataGenerator(dirs, image_size=(64,64), batch_size=batch_size, shuffle=True, max_load=10000, images_only=True )
num_batches = len(gen)
reset_graph()
vae = ConvVAE(z_size=z_size,
batch_size=batch_size,
learning_rate=learning_rate,
kl_tolerance=kl_tolerance,
is_training=True,
reuse=False,
gpu_mode=True,
optimizer=optimizer)
early = EarlyStopping(monitor='loss', min_delta=0.1, patience=5, verbose=verbose, mode='auto')
early.set_model(vae)
early.on_train_begin()
best_loss = sys.maxsize
if verbose:
print("epoch\tstep\tloss\trecon_loss\tkl_loss")
for epoch in range(epochs):
for idx in range(num_batches):
batch = gen[idx]
obs = batch.astype(np.float)/255.0
feed = {vae.x: obs,}
(train_loss, r_loss, kl_loss, train_step, _) = vae.sess.run([
vae.loss, vae.r_loss, vae.kl_loss, vae.global_step, vae.train_op
], feed)
if train_loss < best_loss:
best_loss = train_loss
if save_model:
if ((train_step+1) % 5000 == 0):
vae.save_json("tf_vae/vae.json")
if verbose:
print("{} of {}\t{}\t{:.2f}\t{:.2f}\t{:.2f}".format( epoch, epochs, (train_step+1), train_loss, r_loss, kl_loss) )
gen.on_epoch_end()
early.on_epoch_end(epoch, logs={"loss": train_loss})
if vae.stop_training:
break
early.on_train_end()
# finished, final model:
if save_model:
vae.save_json("tf_vae/vae.json")
return best_loss
gpu_mode=True)
print('Training beta-VAE with beta={:.1f}'.format(beta))
# train loop:
print("train", "step", "loss", "recon_loss", "kl_loss")
for epoch in range(NUM_EPOCH):
np.random.shuffle(dataset)
for idx in range(num_batches):
batch = dataset[idx * batch_size:(idx + 1) * batch_size]
obs = batch.astype(np.float) / 255.0
feed = {
vae.x: obs,
}
(train_loss, r_loss, kl_loss, train_step, _) = vae.sess.run(
[vae.loss, vae.r_loss, vae.kl_loss, vae.global_step, vae.train_op],
feed)
if ((train_step + 1) % 500 == 0):
print("step", (train_step + 1), train_loss, r_loss, kl_loss)
if ((train_step + 1) % 5000 == 0):
vae.save_json(
os.path.join(model_save_path,
args.name + '_vae.json'.format(beta)))
# finished, final model:
vae.save_json(
os.path.join(model_save_path, args.name + '_vae.json'.format(beta)))
# train loop:
print("train", "step", "loss", "recon_loss", "kl_loss")
for epoch in range(arglist.epoch):
np.random.shuffle(dataset)
for idx in range(num_batches):
batch = dataset[idx * arglist.batch_size:(idx + 1) *
arglist.batch_size]
obs = np.array(batch).astype(np.float) / 255.0
feed = {
vae.x: obs,
}
(train_loss, r_loss, kl_loss, train_step, _) = vae.sess.run([
vae.loss, vae.r_loss, vae.kl_loss, vae.global_step,
vae.train_op
], feed)
print("step", (train_step + 1), train_loss, r_loss, kl_loss)
# log.write('{}\n'.format(','.join(map(str,
# [train_loss, r_loss, kl_loss, train_step]))))
if ((train_step + 1) % arglist.save_period == 0):
vae.save_json(
os.path.join(arglist.model_save_path,
"vae_{}.json".format(train_step + 1)))
# finished, final model:
vae.save_json(os.path.join(arglist.model_save_path, "vae.json"))
kl_tolerance=kl_tolerance,
is_training=True,
reuse=False,
gpu_mode=True)
# train loop:
print("train", "step", "loss", "recon_loss", "kl_loss")
for epoch in range(NUM_EPOCH):
np.random.shuffle(dataset)
for idx in range(num_batches):
batch = dataset[idx * batch_size:(idx + 1) * batch_size]
obs = batch.astype(np.float) / 255.0
feed = {
vae.x: obs,
}
(train_loss, r_loss, kl_loss, train_step, _) = vae.sess.run(
[vae.loss, vae.r_loss, vae.kl_loss, vae.global_step, vae.train_op],
feed)
if ((train_step + 1) % 500 == 0):
print("step", (train_step + 1), train_loss, r_loss, kl_loss)
if ((train_step + 1) % 5000 == 0):
vae.save_json("vae/vae.json")
# finished, final model:
vae.save_json("vae/vae.json")
print("Finish time: \t", datetime.datetime.now())