def stylize(content_img,
style_img,
base_img=None,
saveto=None,
gif_step=5,
n_iterations=100,
style_weight=1.0,
content_weight=1.0):
"""Stylization w/ the given content and style images.
Follows the approach in Leon Gatys et al.
Parameters
----------
content_img : np.ndarray
Image to use for finding the content features.
style_img : TYPE
Image to use for finding the style features.
base_img : None, optional
Image to use for the base content. Can be noise or an existing image.
If None, the content image will be used.
saveto : str, optional
Name of GIF image to write to, e.g. "stylization.gif"
gif_step : int, optional
Modulo of iterations to save the current stylization.
n_iterations : int, optional
Number of iterations to run for.
style_weight : float, optional
Weighting on the style features.
content_weight : float, optional
Weighting on the content features.
Returns
-------
stylization : np.ndarray
Final iteration of the stylization.
"""
# Preprocess both content and style images
content_img = vgg16.preprocess(content_img, dsize=(224, 224))[np.newaxis]
style_img = vgg16.preprocess(style_img, dsize=(224, 224))[np.newaxis]
if base_img is None:
base_img = content_img
else:
base_img = make_4d(vgg16.preprocess(base_img, dsize=(224, 224)))
# Get Content and Style features
net = vgg16.get_vgg_model()
g = tf.Graph()
with tf.Session(graph=g) as sess:
tf.import_graph_def(net['graph_def'], name='vgg')
names = [op.name for op in g.get_operations()]
x = g.get_tensor_by_name(names[0] + ':0')
content_layer = 'vgg/conv3_2/conv3_2:0'
content_features = g.get_tensor_by_name(content_layer).eval(
feed_dict={
x: content_img,
'vgg/dropout_1/random_uniform:0': [[1.0] * 4096],
'vgg/dropout/random_uniform:0': [[1.0] * 4096]
})
style_layers = [
'vgg/conv1_1/conv1_1:0', 'vgg/conv2_1/conv2_1:0',
'vgg/conv3_1/conv3_1:0', 'vgg/conv4_1/conv4_1:0',
'vgg/conv5_1/conv5_1:0'
]
style_activations = []
for style_i in style_layers:
style_activation_i = g.get_tensor_by_name(style_i).eval(
feed_dict={
x: style_img,
'vgg/dropout_1/random_uniform:0': [[1.0] * 4096],
'vgg/dropout/random_uniform:0': [[1.0] * 4096]
})
style_activations.append(style_activation_i)
style_features = []
for style_activation_i in style_activations:
s_i = np.reshape(style_activation_i,
[-1, style_activation_i.shape[-1]])
gram_matrix = np.matmul(s_i.T, s_i) / s_i.size
style_features.append(gram_matrix.astype(np.float32))
# Optimize both
g = tf.Graph()
with tf.Session(graph=g) as sess:
net_input = tf.Variable(base_img)
tf.import_graph_def(
net['graph_def'], name='vgg', input_map={'images:0': net_input})
content_loss = tf.nn.l2_loss(
(g.get_tensor_by_name(content_layer) - content_features) /
content_features.size)
style_loss = np.float32(0.0)
for style_layer_i, style_gram_i in zip(style_layers, style_features):
layer_i = g.get_tensor_by_name(style_layer_i)
layer_shape = layer_i.get_shape().as_list()
layer_size = layer_shape[1] * layer_shape[2] * layer_shape[3]
layer_flat = tf.reshape(layer_i, [-1, layer_shape[3]])
gram_matrix = tf.matmul(tf.transpose(layer_flat),
layer_flat) / layer_size
style_loss = tf.add(style_loss,
tf.nn.l2_loss((gram_matrix - style_gram_i) /
np.float32(style_gram_i.size)))
loss = content_weight * content_loss + style_weight * style_loss
optimizer = tf.train.AdamOptimizer(0.01).minimize(loss)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
imgs = []
for it_i in range(n_iterations):
_, this_loss, synth = sess.run(
[optimizer, loss, net_input],
feed_dict={
'vgg/dropout_1/random_uniform:0':
np.ones(
g.get_tensor_by_name('vgg/dropout_1/random_uniform:0')
.get_shape().as_list()),
'vgg/dropout/random_uniform:0':
np.ones(
g.get_tensor_by_name('vgg/dropout/random_uniform:0')
.get_shape().as_list())
})
print(
"iteration %d, loss: %f, range: (%f - %f)" %
(it_i, this_loss, np.min(synth), np.max(synth)),
end='\r')
if it_i % gif_step == 0:
imgs.append(np.clip(synth[0], 0, 1))
if saveto is not None:
gif.build_gif(imgs, saveto=saveto)
return np.clip(synth[0], 0, 1)
def test_gif():
files = CELEB()[:5]
imgs = [imread(f) for f in files]
gif.build_gif(imgs, saveto='test.gif')
assert (os.path.exists('test.gif'))
def train_dataset(ds,
A,
B,
C,
T=20,
n_enc=512,
n_z=200,
n_dec=512,
read_n=12,
write_n=12,
batch_size=100,
n_epochs=100):
if ds is None:
ds = CIFAR10(split=[0.8, 0.1, 0.1])
A, B, C = (32, 32, 3)
n_examples = batch_size
zs = np.random.uniform(-1.0, 1.0, [4, n_z]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
# Clip gradients
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
# Add summary variables
tf.summary.scalar(name='cost', tensor=draw['cost'])
tf.summary.scalar(name='loss_z', tensor=draw['loss_z'])
tf.summary.scalar(name='loss_x', tensor=draw['loss_x'])
tf.summary.histogram(
name='recon_t0_histogram', values=draw['canvas'][0])
tf.summary.histogram(
name='recon_t-1_histogram', values=draw['canvas'][-1])
tf.summary.image(
name='recon_t0_image',
tensor=tf.reshape(draw['canvas'][0], (-1, A, B, C)),
max_outputs=2)
tf.summary.image(
name='recon_t-1_image',
tensor=tf.reshape(draw['canvas'][-1], (-1, A, B, C)),
max_outputs=2)
sums = tf.summary.merge_all()
train_writer = tf.summary.FileWriter(logdir='draw/train')
valid_writer = tf.summary.FileWriter(logdir='draw/valid')
# Init
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
saver = tf.train.Saver()
# Fit all training data
batch_i = 0
test_xs = ds.test.images[:n_examples] / 255.0
utils.montage(test_xs.reshape((-1, A, B, C)), 'draw/test_xs.png')
for epoch_i in range(n_epochs):
for batch_xs, _ in ds.train.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
cost, summary = sess.run(
[draw['cost'], sums, train_op],
feed_dict={
draw['x']: batch_xs / 255.0,
draw['noise']: noise
})[0:2]
train_writer.add_summary(summary, batch_i)
print('train cost:', cost)
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [
utils.montage(r.reshape(-1, A, B, C)) for r in recon
]
gif.build_gif(
recon,
cmap='gray',
saveto='draw/manifold_%08d.gif' % batch_i)
saver.save(sess, './draw/draw.ckpt', global_step=batch_i)
batch_i += 1
# Run validation
if batch_i % 1000 == 0:
for batch_xs, _ in ds.valid.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
cost, summary = sess.run(
[draw['cost'], sums],
feed_dict={
draw['x']: batch_xs / 255.0,
draw['noise']: noise
})[0:2]
valid_writer.add_summary(summary, batch_i)
print('valid cost:', cost)
batch_i += 1
def train_input_pipeline(
files,
A, # img_h
B, # img_w
C,
T=20,
n_enc=512,
n_z=256,
n_dec=512,
read_n=15,
write_n=15,
batch_size=64,
n_epochs=1e9,
input_shape=(64, 64, 3)):
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
batch = create_input_pipeline(
files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=(A, B, C),
shape=input_shape)
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
coord = tf.train.Coordinator()
tf.get_default_graph().finalize()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
# Fit all training data
batch_i = 0
epoch_i = 0
n_files = len(files)
test_xs = sess.run(batch).reshape((-1, A * B * C)) / 255.0
utils.montage(test_xs.reshape((-1, A, B, C)), 'test_xs.png')
try:
while not coord.should_stop() and epoch_i < n_epochs:
batch_xs = sess.run(batch) / 255.0
noise = np.random.randn(batch_size, n_z)
lx, lz = sess.run(
[draw['loss_x'], draw['loss_z'], train_op],
feed_dict={
draw['x']: batch_xs.reshape((-1, A * B * C)) / 255.0,
draw['noise']: noise
})[0:2]
print('x:', lx, 'z:', lz)
if batch_i % n_files == 0:
batch_i = 0
epoch_i += 1
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [
utils.montage(r.reshape(-1, A, B, C)) for r in recon
]
gif.build_gif(recon, saveto='manifold_%08d.gif' % batch_i)
plt.close('all')
batch_i += 1
except tf.errors.OutOfRangeError:
print('Done.')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
# Clean up the session.
sess.close()
def test_gif():
files = CELEB()[:5]
imgs = [imread(f) for f in files]
gif.build_gif(imgs, saveto='test.gif')
assert(os.path.exists('test.gif'))
def test_mnist():
A = 28 # img_h
B = 28 # img_w
C = 1
T = 10
n_enc = 256
n_z = 100
n_dec = 256
read_n = 5
write_n = 5
batch_size = 64
mnist = MNIST(split=[0.8, 0.1, 0.1])
n_examples = batch_size
zs = np.random.uniform(-1.0, 1.0, [4, n_z]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
saver = tf.train.Saver()
# Fit all training data
batch_i = 0
n_epochs = 100
test_xs = mnist.test.images[:n_examples]
utils.montage(test_xs.reshape((-1, A, B)), 'test_xs.png')
for epoch_i in range(n_epochs):
for batch_xs, _ in mnist.train.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
lx, lz = sess.run(
[draw['loss_x'], draw['loss_z'], train_op],
feed_dict={draw['x']: batch_xs,
draw['noise']: noise})[0:2]
print('x:', lx, 'z:', lz)
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [utils.montage(r.reshape(-1, A, B)) for r in recon]
gif.build_gif(
recon,
cmap='gray',
saveto='manifold_%08d.gif' % batch_i)
saver.save(sess, './draw.ckpt', global_step=batch_i)
batch_i += 1
def guided_dream(input_img,
guide_img=None,
downsize=False,
layers=[162, 183, 184, 247],
label_i=962,
layer_i=-1,
feature_loss_weight=1.0,
tv_loss_weight=1.0,
l2_loss_weight=1.0,
softmax_loss_weight=1.0,
model='inception',
neuron_i=920,
n_iterations=100,
save_gif=None,
save_images='imgs',
device='/cpu:0',
**kwargs):
"""Deep Dream v2. Use an optional guide image and other techniques.
Parameters
----------
input_img : np.ndarray
Image to apply deep dream to. Should be 3-dimenionsal H x W x C
RGB uint8 or float32.
guide_img : np.ndarray, optional
Optional image to find features at different layers for. Must pass in
a list of layers that you want to find features for. Then the guided
dream will try to match this images features at those layers.
downsize : bool, optional
Whether or not to downsize the image. Only applies to
model=='inception'.
layers : list, optional
A list of layers to find features for in the "guide_img".
label_i : int, optional
Which label to use for the softmax layer. Use the "get_labels" function
to find the index corresponding the object of interest. If None, not
used.
layer_i : int, optional
Which layer to use for finding the gradient. E.g. the softmax layer
for inception is -1, for vgg networks it is -2. Use the function
"get_layer_names" to find the layer number that you need.
feature_loss_weight : float, optional
Weighting for the feature loss from the guide_img.
tv_loss_weight : float, optional
Total variational loss weighting. Enforces smoothness.
l2_loss_weight : float, optional
L2 loss weighting. Enforces smaller values and reduces saturation.
softmax_loss_weight : float, optional
Softmax loss weighting. Must set label_i.
model : str, optional
Which model to load. Must be one of: ['inception'], 'i2v_tag', 'i2v',
'vgg16', or 'vgg_face'.
neuron_i : int, optional
Which neuron to use. -1 for the entire layer.
n_iterations : int, optional
Number of iterations to dream.
save_gif : bool, optional
Save a GIF.
save_images : str, optional
Folder to save images to.
device : str, optional
Which device to use, e.g. ['/cpu:0'] or '/gpu:0'.
**kwargs : dict
See "_apply" for additional parameters.
Returns
-------
imgs : list of np.ndarray
Images of the dream.
"""
net, img, preprocess, deprocess = _setup(input_img, model, downsize)
print(img.shape, input_img.shape)
print(img.min(), img.max())
if guide_img is not None:
guide_img = preprocess(guide_img.copy(), model)[np.newaxis]
assert (guide_img.shape == img.shape)
batch, height, width, *ch = img.shape
g = tf.Graph()
with tf.Session(graph=g) as sess, g.device(device):
tf.import_graph_def(net['graph_def'], name='net')
names = [op.name for op in g.get_operations()]
input_name = names[0] + ':0'
x = g.get_tensor_by_name(input_name)
features = [names[layer_i] + ':0' for layer_i in layers]
feature_loss = tf.Variable(0.0)
for feature_i in features:
layer = g.get_tensor_by_name(feature_i)
if guide_img is None:
feature_loss += tf.reduce_mean(layer)
else:
# Reshape it to 2D vector
layer = tf.reshape(layer, [-1, 1])
# Do the same for our guide image
guide_layer = sess.run(layer, feed_dict={x: guide_img})
guide_layer = guide_layer.reshape(-1, 1)
# Now calculate their dot product
correlation = tf.matmul(guide_layer.T, layer)
feature_loss += feature_loss_weight * tf.reduce_mean(
correlation)
softmax_loss = tf.Variable(0.0)
if label_i is not None:
layer = g.get_tensor_by_name(names[layer_i] + ':0')
layer_shape = sess.run(tf.shape(layer), feed_dict={x: img})
layer_vec = np.ones(layer_shape) / layer_shape[-1]
layer_vec[..., neuron_i] = 1.0 - 1.0 / layer_shape[1]
softmax_loss += softmax_loss_weight * tf.reduce_mean(
tf.nn.l2_loss(layer - layer_vec))
dx = tf.square(x[:, :height - 1, :width - 1, :] -
x[:, :height - 1, 1:, :])
dy = tf.square(x[:, :height - 1, :width - 1, :] -
x[:, 1:, :width - 1, :])
tv_loss = tv_loss_weight * tf.reduce_mean(tf.pow(dx + dy, 1.2))
l2_loss = l2_loss_weight * tf.reduce_mean(tf.nn.l2_loss(x))
ascent = tf.gradients(feature_loss + softmax_loss + tv_loss + l2_loss,
x)[0]
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
imgs = []
for it_i in range(n_iterations):
this_res, this_feature_loss, this_softmax_loss, this_tv_loss, this_l2_loss = sess.run(
[ascent, feature_loss, softmax_loss, tv_loss, l2_loss],
feed_dict={x: img})
print('feature:', this_feature_loss, 'softmax:', this_softmax_loss,
'tv', this_tv_loss, 'l2', this_l2_loss)
_apply(img, -this_res, it_i, **kwargs)
imgs.append(deprocess(img[0]))
if save_images is not None:
imsave(
os.path.join(save_images, 'frame{}.png'.format(it_i)),
imgs[-1])
if save_gif is not None:
gif.build_gif(imgs, saveto=save_gif)
return imgs
def deep_dream(input_img,
downsize=False,
model='inception',
layer_i=-1,
neuron_i=-1,
n_iterations=100,
save_gif=None,
save_images='imgs',
device='/cpu:0',
**kwargs):
"""Deep Dream with the given parameters.
Parameters
----------
input_img : np.ndarray
Image to apply deep dream to. Should be 3-dimenionsal H x W x C
RGB uint8 or float32.
downsize : bool, optional
Whether or not to downsize the image. Only applies to
model=='inception'.
model : str, optional
Which model to load. Must be one of: ['inception'], 'i2v_tag', 'i2v',
'vgg16', or 'vgg_face'.
layer_i : int, optional
Which layer to use for finding the gradient. E.g. the softmax layer
for inception is -1, for vgg networks it is -2. Use the function
"get_layer_names" to find the layer number that you need.
neuron_i : int, optional
Which neuron to use. -1 for the entire layer.
n_iterations : int, optional
Number of iterations to dream.
save_gif : bool, optional
Save a GIF.
save_images : str, optional
Folder to save images to.
device : str, optional
Which device to use, e.g. ['/cpu:0'] or '/gpu:0'.
**kwargs : dict
See "_apply" for additional parameters.
Returns
-------
imgs : list of np.array
Images of every iteration
"""
net, img, preprocess, deprocess = _setup(input_img, model, downsize)
batch, height, width, *ch = img.shape
g = tf.Graph()
with tf.Session(graph=g) as sess, g.device(device):
tf.import_graph_def(net['graph_def'], name='net')
names = [op.name for op in g.get_operations()]
input_name = names[0] + ':0'
x = g.get_tensor_by_name(input_name)
layer = g.get_tensor_by_name(names[layer_i] + ':0')
layer_shape = sess.run(tf.shape(layer), feed_dict={x: img})
layer_vec = np.ones(layer_shape) / layer_shape[-1]
layer_vec[..., neuron_i] = 1.0 - (1.0 / layer_shape[-1])
ascent = tf.gradients(layer, x)
imgs = []
for it_i in range(n_iterations):
print(it_i, np.min(img), np.max(img))
if neuron_i == -1:
this_res = sess.run(ascent, feed_dict={x: img})[0]
else:
this_res = sess.run(
ascent, feed_dict={x: img,
layer: layer_vec})[0]
_apply(img, this_res, it_i, **kwargs)
imgs.append(deprocess(img[0]))
if save_images is not None:
if not os.path.exists(save_images):
os.mkdir(save_images)
imsave(
os.path.join(save_images, 'frame{}.png'.format(it_i)),
imgs[-1])
if save_gif is not None:
gif.build_gif(imgs, saveto=save_gif)
return imgs
def train_input_pipeline(
files,
A, # img_h
B, # img_w
C,
T=20,
n_enc=512,
n_z=256,
n_dec=512,
read_n=15,
write_n=15,
batch_size=64,
n_epochs=1e9,
input_shape=(64, 64, 3)):
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
batch = create_input_pipeline(
files=files,
batch_size=batch_size,
n_epochs=n_epochs,
crop_shape=(A, B, C),
shape=input_shape)
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
coord = tf.train.Coordinator()
tf.get_default_graph().finalize()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
# Fit all training data
batch_i = 0
epoch_i = 0
n_files = len(files)
test_xs = sess.run(batch).reshape((-1, A * B * C)) / 255.0
utils.montage(test_xs.reshape((-1, A, B, C)), 'test_xs.png')
try:
while not coord.should_stop() and epoch_i < n_epochs:
batch_xs = sess.run(batch) / 255.0
noise = np.random.randn(batch_size, n_z)
lx, lz = sess.run(
[draw['loss_x'], draw['loss_z'], train_op],
feed_dict={
draw['x']: batch_xs.reshape((-1, A * B * C)) / 255.0,
draw['noise']: noise
})[0:2]
print('x:', lx, 'z:', lz)
if batch_i % n_files == 0:
batch_i = 0
epoch_i += 1
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [
utils.montage(r.reshape(-1, A, B, C)) for r in recon
]
gif.build_gif(recon, saveto='manifold_%08d.gif' % batch_i)
plt.close('all')
batch_i += 1
except tf.errors.OutOfRangeError:
print('Done.')
finally:
# One of the threads has issued an exception. So let's tell all the
# threads to shutdown.
coord.request_stop()
# Wait until all threads have finished.
coord.join(threads)
# Clean up the session.
sess.close()
def train_dataset(ds,
A,
B,
C,
T=20,
n_enc=512,
n_z=200,
n_dec=512,
read_n=12,
write_n=12,
batch_size=100,
n_epochs=100):
if ds is None:
ds = CIFAR10(split=[0.8, 0.1, 0.1])
A, B, C = (32, 32, 3)
n_examples = batch_size
zs = np.random.uniform(-1.0, 1.0, [4, n_z]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
# Clip gradients
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
# Add summary variables
tf.summary.scalar(name='cost', tensor=draw['cost'])
tf.summary.scalar(name='loss_z', tensor=draw['loss_z'])
tf.summary.scalar(name='loss_x', tensor=draw['loss_x'])
tf.summary.histogram(
name='recon_t0_histogram', values=draw['canvas'][0])
tf.summary.histogram(
name='recon_t-1_histogram', values=draw['canvas'][-1])
tf.summary.image(
name='recon_t0_image',
tensor=tf.reshape(draw['canvas'][0], (-1, A, B, C)),
max_outputs=2)
tf.summary.image(
name='recon_t-1_image',
tensor=tf.reshape(draw['canvas'][-1], (-1, A, B, C)),
max_outputs=2)
sums = tf.summary.merge_all()
train_writer = tf.summary.FileWriter(logdir='draw/train')
valid_writer = tf.summary.FileWriter(logdir='draw/valid')
# Init
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
saver = tf.train.Saver()
# Fit all training data
batch_i = 0
test_xs = ds.test.images[:n_examples] / 255.0
utils.montage(test_xs.reshape((-1, A, B, C)), 'draw/test_xs.png')
for epoch_i in range(n_epochs):
for batch_xs, _ in ds.train.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
cost, summary = sess.run(
[draw['cost'], sums, train_op],
feed_dict={
draw['x']: batch_xs / 255.0,
draw['noise']: noise
})[0:2]
train_writer.add_summary(summary, batch_i)
print('train cost:', cost)
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [
utils.montage(r.reshape(-1, A, B, C)) for r in recon
]
gif.build_gif(
recon,
cmap='gray',
saveto='draw/manifold_%08d.gif' % batch_i)
saver.save(sess, './draw/draw.ckpt', global_step=batch_i)
batch_i += 1
# Run validation
if batch_i % 1000 == 0:
for batch_xs, _ in ds.valid.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
cost, summary = sess.run(
[draw['cost'], sums],
feed_dict={
draw['x']: batch_xs / 255.0,
draw['noise']: noise
})[0:2]
valid_writer.add_summary(summary, batch_i)
print('valid cost:', cost)
batch_i += 1
def test_mnist():
A = 28 # img_h
B = 28 # img_w
C = 1
T = 10
n_enc = 256
n_z = 100
n_dec = 256
read_n = 5
write_n = 5
batch_size = 64
mnist = MNIST(split=[0.8, 0.1, 0.1])
n_examples = batch_size
zs = np.random.uniform(-1.0, 1.0, [4, n_z]).astype(np.float32)
zs = utils.make_latent_manifold(zs, n_examples)
# We create a session to use the graph
g = tf.Graph()
with tf.Session(graph=g) as sess:
draw = create_model(
A=A,
B=B,
C=C,
T=T,
batch_size=batch_size,
n_enc=n_enc,
n_z=n_z,
n_dec=n_dec,
read_n=read_n,
write_n=write_n)
opt = tf.train.AdamOptimizer(learning_rate=0.0001)
grads = opt.compute_gradients(draw['cost'])
for i, (g, v) in enumerate(grads):
if g is not None:
grads[i] = (tf.clip_by_norm(g, 5), v)
train_op = opt.apply_gradients(grads)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
saver = tf.train.Saver()
# Fit all training data
batch_i = 0
n_epochs = 100
test_xs = mnist.test.images[:n_examples]
utils.montage(test_xs.reshape((-1, A, B)), 'test_xs.png')
for epoch_i in range(n_epochs):
for batch_xs, _ in mnist.train.next_batch(batch_size):
noise = np.random.randn(batch_size, n_z)
lx, lz = sess.run(
[draw['loss_x'], draw['loss_z'], train_op],
feed_dict={draw['x']: batch_xs,
draw['noise']: noise})[0:2]
print('x:', lx, 'z:', lz)
if batch_i % 1000 == 0:
# Plot example reconstructions
recon = sess.run(
draw['canvas'],
feed_dict={draw['x']: test_xs,
draw['noise']: noise})
recon = [utils.montage(r.reshape(-1, A, B)) for r in recon]
gif.build_gif(
recon,
cmap='gray',
saveto='manifold_%08d.gif' % batch_i)
saver.save(sess, './draw.ckpt', global_step=batch_i)
batch_i += 1