VAE (basics, MNIST)

##Library
```python
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

import tensorflow as tf

np.random.seed(10)
tf.set_random_seed(10)
```

##Data
```python
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets('MNIST_data')

images_train = mnist.train.images
labels_train = mnist.train.labels
images_test = mnist.test.images
labels_test = mnist.test.labels

print ('shape: images of train: ', images_train.shape)
print ('shape: labels of train: ', labels_train.shape)
print ('shape: images of test: ', images_test.shape)
print ('shape: leabels of test: ', labels_test.shape)
```
![image.png](https://qiita-image-store.s3.amazonaws.com/0/254604/231318da-61f1-9dc8-277d-fa24c7e68971.png)


##Model
```python
sess.close()
tf.reset_default_graph()

batch_size = 64

X_in = tf.placeholder(dtype=tf.float32, shape=[None, 28, 28])
Y    = tf.placeholder(dtype=tf.float32, shape=[None, 28, 28])
keep_prob = tf.placeholder(dtype=tf.float32, shape=())

Y_flat = tf.reshape(Y, shape=[-1, 28 * 28])

dec_in_channels = 1
n_latent = 8

reshaped_dim = [-1, 7, 7, dec_in_channels]
#inputs_decoder = 49 * dec_in_channels / 2
inputs_decoder = 24
```

```python
def lrelu(x, alpha=0.3):
  return tf.maximum(x, tf.multiply(x, alpha))

def encoder(X_in, keep_prob):
  activation = lrelu
  with tf.variable_scope('encoder', reuse=None):
    X = tf.reshape(X_in, shape=[-1, 28, 28, 1])
    x = tf.layers.conv2d(X, filters=64, kernel_size=4, strides=2, padding='same',
                         activation=activation)
    x = tf.nn.dropout(x, keep_prob)
    x = tf.layers.conv2d(x, filters=64, kernel_size=4, strides=2, padding='same', 
                         activation=activation)
    x = tf.nn.dropout(x, keep_prob)
    x = tf.layers.conv2d(x, filters=64, kernel_size=4, strides=1, padding='same', 
                         activation=activation)
    x = tf.nn.dropout(x, keep_prob)
    x = tf.contrib.layers.flatten(x)
    mn = tf.layers.dense(x, units=n_latent)
    sd = 0.5 * tf.layers.dense(x, units=n_latent)            
    epsilon = tf.random_normal(tf.stack([tf.shape(x)[0], n_latent])) 
    z = mn + tf.multiply(epsilon, tf.exp(sd))

    return z, mn, sd

def decoder(sampled_z, keep_prob):
  with tf.variable_scope('decoder', reuse=None):
    x = tf.layers.dense(sampled_z, units=inputs_decoder, activation=lrelu)
    x = tf.layers.dense(x, units=inputs_decoder * 2 + 1, activation=lrelu)
    x = tf.reshape(x, reshaped_dim)
    x = tf.layers.conv2d_transpose(x, filters=64, kernel_size=4, strides=2, 
                                   padding='same', activation=tf.nn.relu)
    x = tf.nn.dropout(x, keep_prob)
    x = tf.layers.conv2d_transpose(x, filters=64, kernel_size=4, strides=1, 
                                   padding='same', activation=tf.nn.relu)
    x = tf.nn.dropout(x, keep_prob)
    x = tf.layers.conv2d_transpose(x, filters=64, kernel_size=4, strides=1, 
                                   padding='same', activation=tf.nn.relu)

    x = tf.contrib.layers.flatten(x)
    x = tf.layers.dense(x, units=28*28, activation=tf.nn.sigmoid)
    img = tf.reshape(x, shape=[-1, 28, 28])
    return img
```

```python
sampled, mn, sd = encoder(X_in, keep_prob)
dec = decoder(sampled, keep_prob)

unreshaped = tf.reshape(dec, [-1, 28*28])

img_loss = tf.reduce_mean(tf.reduce_sum(tf.squared_difference(unreshaped, Y_flat), 1))
latent_loss = tf.reduce_mean(-0.5 * tf.reduce_sum(1.0 + 2.0 * sd - tf.square(mn) - tf.exp(2.0 * sd), 1))
loss = img_loss + latent_loss

optimizer = tf.train.AdamOptimizer(0.0005)
train_op = optimizer.minimize(loss)

sess = tf.Session()
sess.run(tf.global_variables_initializer())

n_iter = 10000
show_step  = 2000

history_loss = []
history_img_loss = []
history_latent_loss = []

for i in range(n_iter):
  batch = [np.reshape(b, [28, 28]) for b 
           in mnist.train.next_batch(batch_size=batch_size)[0]]
  
  feed_dict = {X_in: batch, Y: batch, keep_prob: 0.8}
  sess.run(train_op, feed_dict=feed_dict )
  temp_loss, temp_img_loss, temp_latent_loss = sess.run([loss, img_loss, latent_loss], 
                                        feed_dict=feed_dict)
  
  history_loss.append(temp_loss)
  history_img_loss.append(temp_img_loss)
  history_latent_loss.append(temp_latent_loss)  

  if not i % show_step:
    ls, d, i_ls, d_ls, mu, sigm = sess.run([loss, dec, img_loss, 
                                            latent_loss, mn, sd], feed_dict = 
                                           {X_in: batch, Y: batch, keep_prob: 1.0})
    print('i: ' +str(i) + ' total loss: ' + str(ls) + ' image loss: ' 
          + str(np.mean(i_ls)) + ' latent loss: ' + str(np.mean(d_ls)))
    
    fig = plt.figure(figsize = (5, 3))
    ax1 = fig.add_subplot(1, 2, 1)
    ax1.imshow(np.reshape(batch[0], [28, 28]), cmap='gray')
    ax1.set_title('Input image')
    ax1.set_axis_off()

    ax1 = fig.add_subplot(1, 2, 2)
    ax1.imshow(np.reshape(d[0], [28, 28]), cmap='gray')
    ax1.set_title('Reconstructed image')
    ax1.set_axis_off()
    
    plt.show()

sess.close()
```
![image.png](https://qiita-image-store.s3.amazonaws.com/0/254604/c9bef282-1328-2d81-5859-e731c110b570.png)

![image.png](https://qiita-image-store.s3.amazonaws.com/0/254604/d0bc8fb5-9ec8-d925-b39d-8d2db7219c31.png)

![image.png](https://qiita-image-store.s3.amazonaws.com/0/254604/41e04c57-ea7c-3ca3-e7f3-2b09ec0e096a.png)

```python
plt.figure(figsize = (5, 3))
plt.plot(range(n_iter), history_loss, 'b-', label='Total')
#plt.plot(range(n_iter), history_img_loss, 'r-', label='Image')
#plt.plot(range(n_iter), history_latent_loss, 'g-', label='Latent')
plt.title('Total loss')
plt.legend(loc='best')
plt.show()

plt.figure(figsize = (5, 3))
#plt.plot(range(n_iter), history_loss, 'b-', label='Total')
plt.plot(range(n_iter), history_img_loss, 'r-', label='Image')
#plt.plot(range(n_iter), history_latent_loss, 'g-', label='Latent')
plt.title('Image (reconstruction) loss')
plt.legend(loc='best')
plt.show()

plt.figure(figsize = (5, 3))
#plt.plot(range(n_iter), history_loss, 'b-', label='Total')
#plt.plot(range(n_iter), history_img_loss, 'r-', label='Image')
plt.plot(range(n_iter), history_latent_loss, 'g-', label='Latent')
plt.title('Latent (kl divergence) loss')
plt.legend(loc='best')
plt.show()
```
![image.png](https://qiita-image-store.s3.amazonaws.com/0/254604/da4f5291-716d-297a-8627-af57a9aa66b6.png)

```python
randoms = [np.random.normal(0, 1, n_latent) for _ in range(10)]
imgs = sess.run(dec, feed_dict = {sampled: randoms, keep_prob: 1.0})
imgs = [np.reshape(imgs[i], [28, 28]) for i in range(len(imgs))]

fig = plt.figure(figsize = (10, 3))

for i, img in enumerate(imgs):
  ax = fig.add_subplot(1, 10, i+1)
  ax.imshow(img, cmap='gray')
  ax.set_axis_off()

plt.show()
```

![image.png](https://qiita-image-store.s3.amazonaws.com/0/254604/466485d3-b366-0360-6b2a-2e9c103a4516.png)