def main(_):
"""High level pipeline.
This script performs the training for VAEs.
"""
# Get dataset.
mnist_dataset = read_data_sets('MNIST_data', one_hot=True)
# Build model.
model = VariationalAutoencoder()
# Start training
train(model, mnist_dataset)
# Plot out latent space, for +/- 3 std.
std = 1
x_z = np.linspace(-3 * std, 3 * std, 20)
y_z = np.linspace(-3 * std, 3 * std, 20)
out = np.empty((28 * 20, 28 * 20))
for x_idx, x in enumerate(x_z):
for y_idx, y in enumerate(y_z):
z_mu = np.array([[y, x]])
img = model.generate_samples(z_mu)
out[x_idx * 28:(x_idx + 1) * 28,
y_idx * 28:(y_idx + 1) * 28] = img[0].reshape(28, 28)
plt.imsave('latent_space_vae.png', out, cmap="gray")
def train_vae(network_architecture, learning_rate=0.001, display_step=5):
vae = VariationalAutoencoder(network_architecture,
learning_rate=learning_rate,
batch_size=BATCH_SIZE)
image_batch, label_batch = inputs(TRAIN_FILE,
batch_size=BATCH_SIZE,
num_epochs=N_EPOCHS)
image_batch_placeholder = tf.placeholder(
tf.float32, shape=[BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, N_CHANNELS])
image_batch = tf.reshape(
image_batch, (BATCH_SIZE, IMAGE_SIZE * IMAGE_SIZE * N_CHANNELS))
# Training cycle
for epoch in range(N_EPOCHS):
avg_cost = 0.
total_batch = int(TRAINING_SET_SIZE / BATCH_SIZE)
# Loop over all batches
for i in range(total_batch):
print("...")
batch_xs = vae.sess.run(image_batch)
# Fit training using batch data
cost = vae.partial_fit(batch_xs)
# Compute average loss
avg_cost += cost / TRAINING_SET_SIZE * BATCH_SIZE
# Display logs per epoch step
if epoch % display_step == 0:
print("Epoch:", '%04d' % (epoch + 1), "cost=",
"{:.9f}".format(avg_cost))
return vae
def main(flags, data):
#############################
''' Experiment Parameters '''
#############################
num_batches = flags.num_batches
dim_z = flags.dim_z
epochs = flags.epochs
learning_rate = flags.lr
l2_loss = flags.l2_loss
seed = flags.seed
data_path = flags.datapath
#Neural Networks parameterising p(x|z), q(z|x)
hidden_layers_px = [600, 600]
hidden_layers_qz = [600, 600]
####################
''' Load Dataset '''
####################
#Uses anglpy module from original paper (linked at top) to load the dataset
train_x, train_y, valid_x, valid_y, test_x, test_y\
= data.load_numpy(data_path, binarize_y=True)
x_train, y_train = train_x.T, train_y.T
x_valid, y_valid = valid_x.T, valid_y.T
x_test, y_test = test_x.T, test_y.T
dim_x = x_train.shape[1]
dim_y = y_train.shape[1]
######################################
''' Train Variational Auto-Encoder '''
######################################
VAE = VariationalAutoencoder(dim_x=dim_x,
dim_z=dim_z,
hidden_layers_px=hidden_layers_px,
hidden_layers_qz=hidden_layers_qz,
l2_loss=l2_loss)
#draw_img uses pylab and seaborn to draw images of original vs. reconstruction
#every n iterations (set to 0 to disable)
VAE.train(x=x_train,
x_valid=x_valid,
epochs=epochs,
num_batches=num_batches,
learning_rate=learning_rate,
seed=seed,
stop_iter=30,
print_every=10,
draw_img=0,
save_path=flags.vaemodel)
def train(network_architecture, sess, learning_rate=0.001,
batch_size=FLAGS.BATCH_SIZE, training_epochs=10,
step_display_step=5, epoch_display_step=5):
vae = VariationalAutoencoder(network_architecture, sess=sess,
transfer_fct=tf.nn.tanh, # FIXME: Fix numerical issues instead of just using tanh
learning_rate=learning_rate,
batch_size=batch_size)
try:
loop_start = datetime.datetime.now()
# Training cycle
for epoch in range(training_epochs):
avg_cost = 0.
total_batch = int(n_samples / batch_size)
# Loop over all batches
for i in range(total_batch):
# batch_xs, _ = mnist.train.next_batch(batch_size)
batch_xs = images_batch.eval()
# Fit training using batch data
cost, r_loss_op, l_loss_op = vae.partial_fit(batch_xs)
vae.save()
if i % step_display_step == 0:
loop_end = datetime.datetime.now()
diff = (loop_end - loop_start).total_seconds()
exps = float(step_display_step) / diff
print("Epoch: (" + str(epoch) + "/" + str(training_epochs) + "); i: (" + str(i) + "/" + str(total_batch) + "). Current cost: " + str(cost) + " (exps: " + ("%.2f" % exps) + ")")
loop_start = datetime.datetime.now()
if np.isnan(cost):
print("Cost is nan!! Killing everything.")
sys.exit(1)
# Compute average loss
avg_cost += cost / n_samples * batch_size
# Display logs per epoch step
if epoch % epoch_display_step == 0:
print("Epoch:", '%04d' % (epoch+1),
"cost=", "{:.9f}".format(avg_cost))
except KeyboardInterrupt:
print("ancelling--Stopping training")
return vae
return vae
class TestVariationalAutoencoder(tf.test.TestCase):
def setUp(self):
self.model = VariationalAutoencoder(0.01, 1.0)
def get_batch_images(self, batch_size):
image = np.zeros((4096), dtype=np.float32)
batch_images = [image] * batch_size
return batch_images
def test_prepare_loss(self):
batch_size = 10
# loss variable should be scalar
self.assertEqual((), self.model.loss.get_shape())
def test_prepare_partial_fit(self):
batch_xs = self.get_batch_images(10)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
# loss result should be float
loss = self.model.partial_fit(sess, batch_xs)
self.assertEqual(np.float32, loss.dtype)
def test_transform(self):
batch_xs = self.get_batch_images(10)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
z_mean, z_log_sigma_sq = self.model.transform(sess, batch_xs)
# check shape of latent variables
self.assertEqual((10, 10), z_mean.shape)
self.assertEqual((10, 10), z_log_sigma_sq.shape)
def test_generate(self):
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
# generate with prior
xs = self.model.generate(sess)
self.assertEqual((1, 4096), xs.shape)
# generate with z_mu with batch size 5
z_mu = np.zeros((5, 10), dtype=np.float32)
xs = self.model.generate(sess, z_mu)
self.assertEqual((5, 4096), xs.shape)
def test_reconstruct(self):
batch_xs = self.get_batch_images(10)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
xs = self.model.reconstruct(sess, batch_xs)
# check reconstructed image shape
self.assertEqual((10, 4096), xs.shape)
def __create_m1(self):
from vae import VariationalAutoencoder
from loss import VariationalInference, kl_divergence_normal
self.model = VariationalAutoencoder(self.dims)
def binary_cross_entropy(r, x):
return F.binary_cross_entropy(r, x, size_average=False)
self.objective = VariationalInference(binary_cross_entropy,
kl_divergence_normal)
self.optimizer_m1 = torch.optim.Adam(self.model.parameters(),
lr=self.lr_m1)
def train(network_architecture,
sess,
learning_rate=0.001,
batch_size=FLAGS.BATCH_SIZE,
training_epochs=10,
display_step=5):
vae = VariationalAutoencoder(
network_architecture,
sess=sess,
transfer_fct=tf.nn.
softplus, # FIXME: Fix numerical issues instead of just using tanh
learning_rate=learning_rate,
batch_size=batch_size)
try:
# Training cycle
for epoch in range(training_epochs):
avg_cost = 0.
total_batch = int(n_samples / batch_size)
# Loop over all batches
for i in range(total_batch):
# batch_xs, _ = mnist.train.next_batch(batch_size)
batch_xs = images_batch.eval()
# Fit training using batch data
cost = vae.partial_fit(batch_xs)
if np.isnan(cost):
raise TrainingException("Got cost=nan")
print("Epoch: (" + str(epoch) + "/" + str(training_epochs) +
"); i: (" + str(i) + "/" + str(total_batch) +
"). Current cost: " + str(cost) + "")
# Compute average loss
avg_cost += cost / n_samples * batch_size
# Display logs per epoch step
if epoch % display_step == 0:
print("Epoch:", '%04d' % (epoch + 1), "cost=",
"{:.9f}".format(avg_cost))
except KeyboardInterrupt:
pass
return vae
def encode_dataset(model_path, min_std=0.0):
VAE = VariationalAutoencoder(
dim_x=28 * 28, dim_z=50) #Should be consistent with model being loaded
with VAE.session:
VAE.saver.restore(VAE.session, VAE_model_path)
enc_x_lab_mean, enc_x_lab_var = VAE.encode(x_lab)
enc_x_ulab_mean, enc_x_ulab_var = VAE.encode(x_ulab)
enc_x_valid_mean, enc_x_valid_var = VAE.encode(x_valid)
enc_x_test_mean, enc_x_test_var = VAE.encode(x_test)
id_x_keep = np.std(enc_x_ulab_mean, axis=0) > min_std
print(id_x_keep)
# enc_x_lab_mean, enc_x_lab_var = enc_x_lab_mean[ :, id_x_keep ], enc_x_lab_var[ :, id_x_keep ]
# enc_x_ulab_mean, enc_x_ulab_var = enc_x_ulab_mean[ :, id_x_keep ], enc_x_ulab_var[ :, id_x_keep ]
# enc_x_valid_mean, enc_x_valid_var = enc_x_valid_mean[ :, id_x_keep ], enc_x_valid_var[ :, id_x_keep ]
# enc_x_test_mean, enc_x_test_var = enc_x_test_mean[ :, id_x_keep ], enc_x_test_var[ :, id_x_keep ]
# data_lab = np.hstack( [ enc_x_lab_mean, enc_x_lab_var ] )
# data_ulab = np.hstack( [ enc_x_ulab_mean, enc_x_ulab_var ] )
# data_valid = np.hstack( [enc_x_valid_mean, enc_x_valid_var] )
# data_test = np.hstack( [enc_x_test_mean, enc_x_test_var] )
# print(data_lab.shape)
data_lab = np.hstack([enc_x_lab_mean, enc_x_lab_var])
data_ulab = np.hstack([enc_x_ulab_mean, enc_x_ulab_var])
data_valid = np.hstack([enc_x_valid_mean, enc_x_valid_var])
data_test = np.hstack([enc_x_test_mean, enc_x_test_var])
print(data_lab.shape)
return data_lab, data_ulab, data_valid, data_test
def encode_dataset(model_path, x_lab, x_ulab, x_valid, x_test, min_std=0.0):
dim_x = x_lab.shape[1]
VAE = VariationalAutoencoder(
dim_x=dim_x, dim_z=50) #Should be consistent with model being loaded
with VAE.session:
VAE.saver.restore(VAE.session, model_path)
enc_x_lab_mean, enc_x_lab_var = VAE.encode(x_lab)
enc_x_ulab_mean, enc_x_ulab_var = VAE.encode(x_ulab)
enc_x_valid_mean, enc_x_valid_var = VAE.encode(x_valid)
enc_x_test_mean, enc_x_test_var = VAE.encode(x_test)
id_x_keep = np.std(enc_x_ulab_mean, axis=0) > min_std
enc_x_lab_mean, enc_x_lab_var = enc_x_lab_mean[:,
id_x_keep], enc_x_lab_var[:,
id_x_keep]
enc_x_ulab_mean, enc_x_ulab_var = enc_x_ulab_mean[:,
id_x_keep], enc_x_ulab_var[:,
id_x_keep]
enc_x_valid_mean, enc_x_valid_var = enc_x_valid_mean[:,
id_x_keep], enc_x_valid_var[:,
id_x_keep]
enc_x_test_mean, enc_x_test_var = enc_x_test_mean[:,
id_x_keep], enc_x_test_var[:,
id_x_keep]
data_lab = np.hstack([enc_x_lab_mean, enc_x_lab_var])
data_ulab = np.hstack([enc_x_ulab_mean, enc_x_ulab_var])
data_valid = np.hstack([enc_x_valid_mean, enc_x_valid_var])
data_test = np.hstack([enc_x_test_mean, enc_x_test_var])
return data_lab, data_ulab, data_valid, data_test
def train(network_architecture,
learning_rate=0.001,
batch_size=100,
training_epochs=10,
display_step=5):
vae = VariationalAutoencoder(network_architecture,
learning_rate=learning_rate,
batch_size=batch_size)
#vae.restore('./tmp/vae_saved.dat')
for epoch in range(training_epochs):
print(epoch)
avg_cost = 0.
total_batch = int(n_samples / batch_size)
for i in range(total_batch):
batch_xs = mnist.train.next_batch(batch_size)[0]
cost = vae.partial_fit(batch_xs)
avg_cost += cost / n_samples * batch_size
if epoch % display_step == 0:
print("Epoch:", '%04d' % (epoch + 1), "cost=",
"{:.9f}".format(avg_cost))
return vae
def __init__(self, sess, input_size, batch_size=128, learning_rate=1e-4):
self.inputs = tf.placeholder(tf.float32, shape=(None, input_size))
network_architecture = \
dict(n_hidden_recog_1=500, # 1st layer encoder neurons
n_hidden_recog_2=500, # 2nd layer encoder neurons
n_hidden_gener_1=500, # 1st layer decoder neurons
n_hidden_gener_2=500, # 2nd layer decoder neurons
n_input=input_size, # input size
n_z=input_size) # dimensionality of latent space
self.ema = pyma.EMA(0.1)
self.iteration = 0
self.loss_vae = VariationalAutoencoder(sess, network_architecture,
learning_rate=learning_rate,
batch_size=batch_size)
def encode_dataset( model_path, min_std = 0.0 ):
VAE = VariationalAutoencoder( dim_x = 28*28, dim_z = 50 ) #Should be consistent with model being loaded
with VAE.session:
VAE.saver.restore( VAE.session, VAE_model_path )
enc_x_lab_mean, enc_x_lab_var = VAE.encode( x_lab )
enc_x_ulab_mean, enc_x_ulab_var = VAE.encode( x_ulab )
enc_x_valid_mean, enc_x_valid_var = VAE.encode( x_valid )
enc_x_test_mean, enc_x_test_var = VAE.encode( x_test )
id_x_keep = np.std( enc_x_ulab_mean, axis = 0 ) > min_std
enc_x_lab_mean, enc_x_lab_var = enc_x_lab_mean[ :, id_x_keep ], enc_x_lab_var[ :, id_x_keep ]
enc_x_ulab_mean, enc_x_ulab_var = enc_x_ulab_mean[ :, id_x_keep ], enc_x_ulab_var[ :, id_x_keep ]
enc_x_valid_mean, enc_x_valid_var = enc_x_valid_mean[ :, id_x_keep ], enc_x_valid_var[ :, id_x_keep ]
enc_x_test_mean, enc_x_test_var = enc_x_test_mean[ :, id_x_keep ], enc_x_test_var[ :, id_x_keep ]
data_lab = np.hstack( [ enc_x_lab_mean, enc_x_lab_var ] )
data_ulab = np.hstack( [ enc_x_ulab_mean, enc_x_ulab_var ] )
data_valid = np.hstack( [enc_x_valid_mean, enc_x_valid_var] )
data_test = np.hstack( [enc_x_test_mean, enc_x_test_var] )
return data_lab, data_ulab, data_valid, data_test
def train(network_architecture, sess, learning_rate=0.001,
batch_size=100, training_epochs=10, display_step=5):
vae = VariationalAutoencoder(network_architecture, sess=sess,
# transfer_fct=tf.nn.tanh,
learning_rate=learning_rate,
batch_size=batch_size)
# Training cycle
for epoch in range(training_epochs):
avg_cost = 0.
total_batch = int(n_samples / batch_size)
# Loop over all batches
for i in range(total_batch):
batch_xs, _ = mnist.train.next_batch(batch_size)
# Fit training using batch data
cost = vae.partial_fit(batch_xs)
# Compute average loss
avg_cost += cost / n_samples * batch_size
# Display logs per epoch step
if epoch % display_step == 0:
print("Epoch:", '%04d' % (epoch+1),
"cost=", "{:.9f}".format(avg_cost))
return vae
class LossNet:
def __init__(self, sess, input_size, batch_size=128, learning_rate=1e-4):
self.inputs = tf.placeholder(tf.float32, shape=(None, input_size))
network_architecture = \
dict(n_hidden_recog_1=500, # 1st layer encoder neurons
n_hidden_recog_2=500, # 2nd layer encoder neurons
n_hidden_gener_1=500, # 1st layer decoder neurons
n_hidden_gener_2=500, # 2nd layer decoder neurons
n_input=input_size, # input size
n_z=input_size) # dimensionality of latent space
self.ema = pyma.EMA(0.1)
self.iteration = 0
self.loss_vae = VariationalAutoencoder(sess, network_architecture,
learning_rate=learning_rate,
batch_size=batch_size)
def step(self, minibatch):
minibatch = np.asarray(minibatch)
batch_size = minibatch.shape[0]
minibatch_var = np.var(minibatch) + 1e-9
normalized_minibatch = (minibatch - np.mean(minibatch)) / minibatch_var
# normalized_minibatch = preprocessing.scale(minibatch)
_, z_m, z_log_sigma_sq = self.loss_vae.step(normalized_minibatch)
# Update our moving average
ema_update = round(self.ema.compute(np.sum(z_m)), 5)
var_update = 3.0*np.sqrt(np.abs(np.sum(z_log_sigma_sq)))
self.iteration += 1
loss_scaler = np.abs(ema_update) > var_update
loss_scaler = loss_scaler.astype(float)
# if self.iteration > 3000: # settling time
# # If this is over 3x the std deviation then allow info flow
# loss_scaler = ema_update > var_update
# loss_scaler = loss_scaler.astype(float)
# else:
# loss_scaler = 1.0
print ' ema update: ', '{:.2f}'.format(ema_update), \
' | 3*sigma: ', '{:.2f}'.format(var_update), \
' | loss_scaler: ', loss_scaler,
return loss_scaler
def main(argv):
manager = DataManager()
manager.load()
sess = tf.Session()
model = VariationalAutoencoder(learning_rate=flags.learning_rate,
beta=flags.beta)
sess.run(tf.global_variables_initializer())
saver = load_checkpoints(sess)
if flags.training:
# Train
train(sess, model, manager, saver)
else:
reconstruct_check_images = manager.get_random_images(10)
# Image reconstruction check
reconstruct_check(sess, model, reconstruct_check_images)
# Disentangle check
disentangle_check(sess, model, manager)
help="path of the file containing SMILES strings for reconstruction")
parser.add_argument(
"reconst_filename",
help=
"path of file to write reconstructed strings to, for checking accuracy"
)
args = parser.parse_args()
# path of file to write reconstructed strings to for checking accuracy
#reconst_filename = "smiles_small_reconstructed.fp"
# Obtain the dataset for testing
tdata, num_samples, w, h, input_dim = get_data(args.test_filename)
# Test the model for reconstruction
loaded_model = VariationalAutoencoder(input_dim)
loaded_model.restore_model('./model')
print("Loaded model from disk")
reconstructed = test_reconstruction(loaded_model, tdata, num_samples)
#writing to a text file
with open(args.reconst_filename, "w") as File:
for i, r in enumerate(reconstructed):
for item in r:
File.write("%f " % item)
File.write("\n")
File.close()
parser.add_argument('--batch_size',
type=int,
default=100,
help='batch size')
parser.add_argument('--hidden_size',
type=int,
default=200,
help='hidden size of network layers')
args = parser.parse_args()
data_set_name = args.dataset
vae_type = args.vae_type
result_dir = './pictures/new_pictures/'
input_values, width, height, nr_channels = data_reader.read_data_set(
data_set_name)
v2 = VariationalAutoencoder(width,
height,
nr_channels,
dim_latent=2,
input_values=input_values,
type=vae_type,
learning_rate=args.learning_rate,
batch_size=args.batch_size,
dim_hidden=args.hidden_size)
model_path_2 = '../checkpoints/' + data_set_name + '_' + vae_type + '_2'
v2 = restore_model(v2, model_path_2)
plot_latent_space_2d(v2, input_values.test.images,
input_values.test.labels)
def train(network_architecture,
sess,
learning_rate=FLAGS.LEARNING_RATE,
batch_size=FLAGS.BATCH_SIZE,
training_epochs=FLAGS.MAX_EPOCHS,
step_display_step=5,
epoch_display_step=5):
"""
The training function builds the autoencoder and trains it.
:param network_architecture: a set of variables corresponding to values defining
the network architecture. See documentation for details.
:param sess: the tensorflow session.
:param learning_rate: the learning rate for the optimizer.
Defaults to value stored in FLAGS.
:param batch_size: size of the images batch.
Defaults to value stored in FLAGS.
:param training_epochs: Number of epochs to train.
Defaults to value stored in FLAGS.
:param step_display_step: Display training info every {N} steps, default is 5.
:param epoch_display_step: Display training info every {N} epochs, default is 5.
:return:
"""
vae = VariationalAutoencoder(
network_architecture,
sess=sess,
transfer_fct=tf.nn.
tanh, # FIXME: Fix numerical issues instead of just using tanh
learning_rate=learning_rate,
batch_size=batch_size)
try:
loop_start = datetime.datetime.now()
# Training cycle
for epoch in range(training_epochs):
avg_cost = 0.
total_batch = int(n_samples / batch_size)
# Loop over all batches
for i in range(total_batch):
# batch_xs, _ = mnist.train.next_batch(batch_size)
batch_xs = images_batch.eval()
# Fit training using batch data
cost, r_loss_op, l_loss_op = vae.partial_fit(batch_xs)
if i % step_display_step == 0:
loop_end = datetime.datetime.now()
diff = (loop_end - loop_start).total_seconds()
exps = float(step_display_step) / diff
magic_print("Epoch: (" + str(epoch) + "/" +
str(training_epochs) + "); i: (" + str(i) +
"/" + str(total_batch) + "). Current cost: " +
str(cost) + " (exps: " +
("%.2f" % exps) + ")",
epoch=epoch,
max_epochs=training_epochs,
i=i,
max_i=total_batch)
loop_start = datetime.datetime.now()
if np.isnan(cost):
print("Cost is nan!! Killing everything.")
sys.exit(1)
else:
# Only save the state of the autoencoder if the cost is valid.
# (This is needed, because if we save an autoencoder with
# a cost of nan, we can't really load its weights, because
# they're basically useless).
vae.save()
# Compute average loss
avg_cost += cost / n_samples * batch_size
# Display logs per epoch step
if epoch % epoch_display_step == 0:
print("Epoch:", '%04d' % (epoch + 1), "cost=",
"{:.9f}".format(avg_cost))
except KeyboardInterrupt:
print("ancelling--Stopping training")
return vae
return vae
def exp1(overwrite=False):
"""Vanilla autoencoder. Gaussian prior and posterior, Bernoulli likelihood."""
print("\nRunning VAE with two dimensional latency space...\n")
latency_dim = 2
encoder = [500] * 3 # num neurons in each layer of encoder network
decoder = [500] * 3 # num neurons in each layer of generator network
network_architecture = \
dict(n_hidden_recog=encoder, # num neurons in each layer of encoder network
n_hidden_gener=decoder, # num neurons in each layer of generator network
n_input=input_dim, # (img shape: 28*28)
n_z=latency_dim) # dimensionality of latent space
training_epochs = 4
batch_size = 100
learning_rate = 0.0001
filename = 'experiments/latency_dim=%d/enc=%s_dec=%s_epochs=%d_batches=%d_opt=adam_learnRate=%s' \
% (latency_dim, '-'.join([str(l) for l in encoder]),
'-'.join([str(l) for l in decoder]),
training_epochs, batch_size, float('%.4g' % learning_rate))
if not os.path.exists(
os.path.join(filename, 'latency_space_while_training')):
os.makedirs(os.path.join(filename, 'latency_space_while_training'))
if os.path.exists('./' + filename + '/exp.meta') and not overwrite:
vae = VariationalAutoencoder(network_architecture)
new_saver = tf.train.import_meta_graph('./' + filename + '/exp.meta')
new_saver.restore(vae.sess,
tf.train.latest_checkpoint('./' + filename + '/'))
else:
vae = VariationalAutoencoder(network_architecture,
learning_rate=learning_rate,
batch_size=batch_size)
numPlots = 0
nx = ny = 20
rnge = 4
x_values = np.linspace(-rnge, rnge, nx)
y_values = np.linspace(-rnge, rnge, ny)
interval = 10 # save fig of latency space every 10 batches
# Training cycle
for epoch in range(training_epochs):
avg_cost = 0.
avg_cost2 = 0.
total_batch = int(n_samples / batch_size)
# Loop over all batches
canvas = np.empty((winSize * ny, winSize * nx))
for i in range(total_batch):
batch_xs = synthData[i * batch_size:(i + 1) * batch_size]
# Fit training using batch data
cost = vae.partial_fit(batch_xs, batch_xs)
# Compute average loss
avg_cost += cost / n_samples * batch_size
avg_cost2 += cost / (batch_size * interval) * batch_size
if i % interval == 0:
fig = plt.figure()
axes = fig.add_subplot(1, 2, 1)
axes2 = fig.add_subplot(1, 2, 2)
axes.set_aspect('equal')
axes.set_title("%d Iterations" %
(epoch * n_samples + i * batch_size))
axes2.set_title("Average Cost = %f" % (avg_cost2))
avg_cost2 = 0.
axes.set_xlim((-rnge, rnge))
axes.set_ylim((-rnge, rnge))
test_xs = synthData[19000:20000]
test_ys = synthDataLabels[19000:20000]
means, stds = vae.getLatentParams(test_xs)
latent_xs = means + stds * np.random.normal(
size=stds.shape)
im = axes.scatter(latent_xs[:, 0],
latent_xs[:, 1],
c=test_ys,
edgecolor='k')
cax = fig.add_axes([0.09, 0.71, 0.37, 0.03])
fig.colorbar(im, cax=cax, orientation='horizontal')
ws = winSize
for i, yi in enumerate(x_values):
for j, xi in enumerate(y_values):
z_mu = np.array([[xi, yi]] * vae.batch_size)
x_mean = vae.generate(z_mu)
canvas[(nx - i - 1) * ws:(nx - i) * ws,
j * ws:(j + 1) * ws] = x_mean[0].reshape(
ws, ws)
axes2.imshow(canvas, origin="upper")
plt.tight_layout()
fig.savefig(filename +
'/latency_space_while_training/fig%04d.png' %
numPlots)
numPlots += 1
plt.close()
# Display logs per epoch step
print("Epoch:", '%04d' % (epoch), "cost=",
"{:.9f}".format(avg_cost))
vae.saver.save(vae.sess, filename + '/exp.meta')
class ssl_vae:
def __init__(self,
X_labeled,
Y_labeled,
X_unlabeled,
batch_size=10,
num_workers=1,
lr_m1=3e-5,
lr_m2=1e-2,
epochs_m1=50,
epochs_m2=50,
dims=[3, 32, [128]],
verbose=False):
"""
Args:
X_labeled (array): for now accepts a flat tensor [observations x k] [observations x height x weight x channels]
Y_labeled (array): [observations x 1]
X_unlabeled (array): [observations x height x weight x channels]
batch_size (int): size of batch sample
num_workers (int): number of workers on each dataset
lr_m1,lr_m2 (int,int): learning rates for model1 and model2
epochs_m1,epochs_m2 (int,int): number of epochs to train model1 and model2 for
dims (list): has the format [x_dim, z_dim, [h_dim]] , creates the VAE
verbose (boolean): prints the training process of the VAE
"""
#flatten_bernoulli = transforms.Lambda(lambda img: transforms.ToTensor()(img).view(-1).bernoulli())
self.lr_m1 = lr_m1
self.lr_m2 = lr_m2
self.epochs_m1 = epochs_m1
self.epochs_m2 = epochs_m2
self.dims = dims
self.verbose = verbose
labeled = ssl_vae_dataset(X_labeled, Y_labeled)
unlabeled = ssl_vae_dataset(X_labeled)
self.unlabeled = DataLoader(unlabeled,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
self.labeled = DataLoader(labeled,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
def train(self):
self.__create_m1() # create m1
self.__train_m1() # train m1
self.__create_m2() # create m2
self.__train_m2() # train m2
def __create_m1(self):
from vae import VariationalAutoencoder
from loss import VariationalInference, kl_divergence_normal
self.model = VariationalAutoencoder(self.dims)
def binary_cross_entropy(r, x):
return F.binary_cross_entropy(r, x, size_average=False)
self.objective = VariationalInference(binary_cross_entropy,
kl_divergence_normal)
self.optimizer_m1 = torch.optim.Adam(self.model.parameters(),
lr=self.lr_m1)
def __train_m1(self):
for epoch in range(self.epochs_m1):
for u in self.unlabeled:
u = u['X']
u = Variable(u)
reconstruction, (_, z_mu, z_log_var) = self.model(u)
# Equation 3
self.L = self.objective(reconstruction, u, z_mu, z_log_var)
self.L.backward()
self.optimizer_m1.step()
self.optimizer_m1.zero_grad()
if self.verbose and epoch % 10 == 0:
l = self.L.data[0]
print("Epoch: {0:} loss: {1:.3f}".format(epoch, l))
def __create_m2(self):
import torch.nn as nn
self.classifier_m2 = nn.Sequential(nn.Linear(32, 32),
nn.ReLU(inplace=True),
nn.Linear(32, 10), nn.Softmax())
self.optimizer_m2 = torch.optim.Adam(self.classifier_m2.parameters(),
lr=self.lr_m2)
def __train_m2(self):
for epoch in range(self.epochs_m2):
for l in self.labeled:
x = l['X']
y = l['y']
x = Variable(x)
y = Variable(y)
_, (z, _, _) = self.model(x)
logits = self.classifier_m2(z)
self.loss = F.cross_entropy(logits, y)
self.loss.backward()
self.optimizer_m2.step()
self.optimizer_m2.zero_grad()
if self.verbose and epoch % 10 == 0:
l = self.loss.data[0]
print("Epoch: {0:} loss: {1:.3f}".format(epoch, l))
def predict(self, X):
_, (z, _, _) = self.model(Variable(X))
logits = self.classifier_m2(z)
_, prediction = torch.max(logits, 1)
return prediction.data
mnist_path, binarize_y=True)
x_train, y_train = train_x.T, train_y.T
x_valid, y_valid = valid_x.T, valid_y.T
x_test, y_test = test_x.T, test_y.T
dim_x = x_train.shape[1]
dim_y = y_train.shape[1]
######################################
''' Train Variational Auto-Encoder '''
######################################
VAE = VariationalAutoencoder(dim_x=dim_x,
dim_z=dim_z,
hidden_layers_px=hidden_layers_px,
hidden_layers_qz=hidden_layers_qz,
l2_loss=l2_loss)
# draw_img uses pylab and seaborn to draw images of original vs. reconstruction
# every n iterations (set to 0 to disable)
VAE.train(x=x_train,
x_valid=x_valid,
epochs=epochs,
num_batches=num_batches,
learning_rate=learning_rate,
seed=seed,
stop_iter=30,
print_every=10,
draw_img=1)
def setUp(self):
self.model = VariationalAutoencoder(ndims=20, nlatent=2)
import numpy as np
import tensorflow as tf
#import matplotlib.pyplot as plt
#%matplotlib inline
from tensorflow.examples.tutorials.mnist import input_data
#import vae
from vae import VariationalAutoencoder
np.random.seed(66)
tf.set_random_seed(666)
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
n_samples = mnist.train.num_examples
network_architecture = \
dict(n_hidden_recog_1 = 500,
n_hidden_recog_2 = 500,
n_hidden_gener_1 = 500,
n_hidden_gener_2 = 500,
n_input = 784,
n_z = 20)
learning_rate = 0.001
batch_size = 100
vae = VariationalAutoencoder(network_architecture,
learning_rate=learning_rate,
batch_size=batch_size)
vae.restore('./tmp/vae_saved.dat')
def setUp(self):
self.model = VariationalAutoencoder(0.01, 1.0)
