def train(args):
learning_rate = args.learning_rate
batch_size = args.batch_size
training_epochs = args.training_epochs
keep_prob = args.keep_prob
checkpoint_step = args.checkpoint_step # save training results every check point step
z_dim = args.z_dim # number of latent variables.
loss_mode = args.loss_mode
diff_mode = False
if args.diff_mode == 1:
diff_mode = True
dirname = 'save'
if not os.path.exists(dirname):
os.makedirs(dirname)
with open(os.path.join(dirname, 'config.pkl'), 'w') as f:
cPickle.dump(args, f)
vae = VAE(learning_rate=learning_rate, batch_size=batch_size, z_dim = z_dim, keep_prob = keep_prob, loss_mode = loss_mode)
mnist = read_data_sets()
n_samples = mnist.num_examples
# load previously trained model if appilcable
ckpt = tf.train.get_checkpoint_state(dirname)
if ckpt:
vae.load_model(dirname)
# Training cycle
for epoch in range(training_epochs):
avg_cost = 0.
avg_likelihood_loss = 0.
avg_kl_loss = 0.
mnist.shuffle_data()
total_batch = int(n_samples / batch_size)
# Loop over all batches
for i in range(total_batch):
batch_xs = mnist.next_batch(batch_size)
if (diff_mode == True):
batch_xs = mnist.integrate_batch(batch_xs)
# Fit training using batch data
cost, likelihood_loss, kl_loss = vae.partial_fit(batch_xs)
# Compute average loss
avg_cost += cost / n_samples * batch_size
avg_likelihood_loss += likelihood_loss / n_samples * batch_size
avg_likelihood_loss += kl_loss / n_samples * batch_size
# Display logs per batch
'''
print "batch:", '%04d' % (i+1), \
"total loss =", "{:.6f}".format(cost), \
"likelihood_loss =", "{:.6f}".format(likelihood_loss), \
"kl_loss =", "{:.6f}".format(kl_loss)
'''
# Display logs per epoch step
print "Epoch:", '%04d' % (epoch+1), \
"total loss =", "{:.6f}".format(avg_cost), \
"likelihood_loss =", "{:.6f}".format(avg_likelihood_loss), \
"kl_loss =", "{:.6f}".format(avg_kl_loss)
# save model
if epoch > 0 and epoch % checkpoint_step == 0:
checkpoint_path = os.path.join('save', 'model.ckpt')
vae.save_model(checkpoint_path, epoch)
print "model saved to {}".format(checkpoint_path)
# save model one last time, under zero label to denote finish.
vae.save_model(checkpoint_path, 0)
return vae
class Sampler():
def __init__(self):
self.mnist = None
self.model = VAE()
self.model.load_model('save')
self.z = self.generate_z()
def get_random_mnist(self, with_label=False):
if self.mnist == None:
self.mnist = read_data_sets()
if with_label == True:
data, label = self.mnist.next_batch(1, with_label)
return data[0], label[0]
return self.mnist.next_batch(1)[0]
def get_random_specific_mnist(self, label=2):
m, l = self.get_random_mnist(with_label=True)
for i in range(100):
if l == label:
break
m, l = self.get_random_mnist(with_label=True)
return m
def generate_random_label(self, label):
m = self.get_random_specific_mnist(label)
self.show_image(m)
self.show_image_from_z(self.encode(m))
def generate_z(self):
z = np.random.normal(size=(1, self.model.z_dim)).astype(np.float32)
return z
def encode(self, mnist_data):
new_shape = [1] + list(mnist_data.shape)
return self.model.transform(np.reshape(mnist_data, new_shape))
def generate(self, z=None):
if z is None:
z = self.generate_z()
else:
z = np.reshape(z, (1, self.model.z_dim))
self.z = z
return self.model.generate(z)[0]
def show_image(self, image_data):
'''
image_data is a tensor, in [height width depth]
image_data is NOT the PIL.Image class
'''
plt.subplot(1, 1, 1)
y_dim = image_data.shape[0]
x_dim = image_data.shape[1]
c_dim = 1
if c_dim > 1:
plt.imshow(image_data, interpolation='nearest')
else:
plt.imshow(image_data.reshape(y_dim, x_dim),
cmap='Greys',
interpolation='nearest')
plt.axis('off')
plt.show()
def show_image_from_z(self, z):
self.show_image(self.generate(z))
def to_image(self, image_data):
# convert to PIL.Image format from np array (0, 1)
img_data = np.array(1 - image_data)
y_dim = image_data.shape[0]
x_dim = image_data.shape[1]
c_dim = 1
if c_dim > 1:
img_data = np.array(img_data.reshape(
(y_dim, x_dim, c_dim)) * 255.0,
dtype=np.uint8)
else:
img_data = np.array(img_data.reshape((y_dim, x_dim)) * 255.0,
dtype=np.uint8)
im = Image.fromarray(img_data)
return im
def diff_image(self, image_data):
# perform 2d differentiation on mnist image
m2 = np.array(image_data) # makes a copy
m2[1:, 1:, :] = m2[1:, 1:, :] - m2[0:-1, 1:, :]
m2[1:, 1:, :] = m2[1:, 1:, :] - m2[1:, 0:-1, :]
return m2
def integrate_image(self, image_data):
# integrates differentiated batch back to mnist image
m3 = np.array(image_data)
m3 = m3.cumsum(axis=0)
m3 = m3.cumsum(axis=1)
return m3
class Sampler():
def __init__(self):
self.mnist = None
self.model = VAE()
self.model.load_model('save')
self.z = self.generate_z()
def get_random_mnist(self, with_label = False):
if self.mnist == None:
self.mnist = read_data_sets()
if with_label == True:
data, label = self.mnist.next_batch(1, with_label)
return data[0], label[0]
return self.mnist.next_batch(1)[0]
def get_random_specific_mnist(self, label = 2):
m, l = self.get_random_mnist(with_label = True)
for i in range(100):
if l == label:
break
m, l = self.get_random_mnist(with_label = True)
return m
def generate_random_label(self, label):
m = self.get_random_specific_mnist(label)
self.show_image(m)
self.show_image_from_z(self.encode(m))
def generate_z(self):
z = np.random.normal(size=(1, self.model.z_dim)).astype(np.float32)
return z
def encode(self, mnist_data):
new_shape = [1]+list(mnist_data.shape)
return self.model.transform(np.reshape(mnist_data, new_shape))
def generate(self, z=None):
if z is None:
z = self.generate_z()
else:
z = np.reshape(z, (1, self.model.z_dim))
self.z = z
return self.model.generate(z)[0]
def show_image(self, image_data):
'''
image_data is a tensor, in [height width depth]
image_data is NOT the PIL.Image class
'''
plt.subplot(1, 1, 1)
y_dim = image_data.shape[0]
x_dim = image_data.shape[1]
c_dim = 1
if c_dim > 1:
plt.imshow(image_data, interpolation='nearest')
else:
plt.imshow(image_data.reshape(y_dim, x_dim), cmap='Greys', interpolation='nearest')
plt.axis('off')
plt.show()
def show_image_from_z(self, z):
self.show_image(self.generate(z))
def to_image(self, image_data):
# convert to PIL.Image format from np array (0, 1)
img_data = np.array(1-image_data)
y_dim = image_data.shape[0]
x_dim = image_data.shape[1]
c_dim = 1
if c_dim > 1:
img_data = np.array(img_data.reshape((y_dim, x_dim, c_dim))*255.0, dtype=np.uint8)
else:
img_data = np.array(img_data.reshape((y_dim, x_dim))*255.0, dtype=np.uint8)
im = Image.fromarray(img_data)
return im
def diff_image(self, image_data):
# perform 2d differentiation on mnist image
m2 = np.array(image_data) # makes a copy
m2[1:,1:,:] = m2[1:,1:,:]-m2[0:-1,1:,:]
m2[1:,1:,:] = m2[1:,1:,:]-m2[1:,0:-1,:]
return m2
def integrate_image(self, image_data):
# integrates differentiated batch back to mnist image
m3 = np.array(image_data)
m3 = m3.cumsum(axis=0)
m3 = m3.cumsum(axis=1)
return m3
