Structured Output Prediction (Conditional Variational Autoencoder) の実装に関するメモ

##ポイント
* Conditional Variational Autoencoder を実装し、具体的な数値で確認。

##レファレンス
[1. Learning Structured Output Representation using Deep Conditional Generative Models](https://pdfs.semanticscholar.org/3f25/e17eb717e5894e0404ea634451332f85d287.pdf)

##サンプルコード
```python
# Gaussian
class ConditionalVariationalAutoencoder():
  
  def __init__(self):
    pass
  
  def weight_variable(self, name, shape):
    initializer = tf.truncated_normal_initializer(mean = 0.0, stddev = 0.01, dtype = tf.float32)
    return tf.get_variable(name, shape, initializer = initializer)
  
  def bias_variable(self, name, shape):
    initializer = tf.constant_initializer(value = 0.0, dtype = tf.float32)
    return tf.get_variable(name, shape, initializer = initializer)

  def beta_variable(self, name, shape):
    initializer = tf.constant_initializer(value = 0.0, dtype = tf.float32)
    return tf.get_variable(name, shape, initializer = initializer)

  def gamma_variable(self, name, shape):
    initializer = tf.constant_initializer(value = 1.0, dtype = tf.float32)
    return tf.get_variable(name, shape, initializer = initializer)

  def kl_divergence_2(self, mu_q, logvar_q, mu_p, logvar_p, dim):
    sig_q = tf.exp(0.5 * logvar_q)
    sig_p = tf.exp(0.5 * logvar_p)
    calc_1 = tf.reduce_sum((sig_q * sig_q) / (sig_p * sig_p), axis = 1)
    calc_2 = tf.reduce_sum(tf.square(mu_q - mu_p) / (sig_p * sig_p), axis = 1)
    calc_3 = 2.0 * tf.reduce_sum(tf.log(sig_q / sig_p), axis = 1)

    return 0.5 * (calc_1 + calc_2 - dim - calc_3)

  def kl_divergence(self, mu_q, logvar_q, mu_p, logvar_p):
    sig_q = tf.exp(0.5 * logvar_q)
    sig_p = tf.exp(0.5 * logvar_p)
    dist_q = tf.contrib.distributions.MultivariateNormalDiag(loc = mu_q, scale_diag = sig_q)
    dist_p = tf.contrib.distributions.MultivariateNormalDiag(loc = mu_p, scale_diag = sig_p)

    #return tf.contrib.distributions.kl_divergence(dist_q, dist_p)
    return dist_q.kl_divergence(dist_p)
  
  def inference(self, x, y, dim_x, dim_y, n_units_1, dim_z, batch_size):
    with tf.variable_scope('enc_z_dec'):
     #for z_xy and y_xz
      w_1 = self.weight_variable('w_1', [dim_x + dim_y, n_units_1])
      b_1 = self.bias_variable('b_1', [n_units_1])

      xy = tf.concat([x, y], axis = 1)
      h_1 = tf.add(tf.matmul(xy, w_1), b_1)
      
      # batch normalization
      beta_1 = self.beta_variable('beta_1', [n_units_1])
      gamma_1 = self.gamma_variable('gamma_1', [n_units_1])
      
      mean_1, var_1 = tf.nn.moments(h_1, [0])
      h_1 = gamma_1 * (h_1 - mean_1) / tf.sqrt(var_1 + 1e-5) + beta_1
      
      bn_1 = tf.nn.softplus(h_1)

      w_2 = self.weight_variable('w_2', [n_units_1, dim_z * 2])
      b_2 = self.bias_variable('b_2', [dim_z * 2])

      h_2 = tf.add(tf.matmul(bn_1, w_2), b_2)
      
      mu_z_xy = h_2[:, : dim_z]
      logvar_z_xy = h_2[:, dim_z :]
            
      z_xy = mu_z_xy + tf.random_normal(shape = [batch_size, dim_z], mean = 0.0, stddev = 1.0) * tf.exp(0.5 * logvar_z_xy)
      
      w_3 = self.weight_variable('w_3', [dim_x + dim_z, n_units_1])
      b_3 = self.bias_variable('b_3', [n_units_1])

      x_z_xy = tf.concat([x, z_xy], axis = 1)
      h_3 = tf.add(tf.matmul(x_z_xy, w_3), b_3)

      # batch normalization
      beta_2 = self.beta_variable('beta_2', [n_units_1])
      gamma_2 = self.gamma_variable('gamma_2', [n_units_1])
      
      mean_2, var_2 = tf.nn.moments(h_3, [0])
      h_3 = gamma_2 * (h_3 - mean_2) / tf.sqrt(var_2 + 1e-5) + beta_2
      
      bn_2 = tf.nn.softplus(h_3)
      
      w_4 = self.weight_variable('w_4', [n_units_1, dim_y])
      b_4 = self.bias_variable('b_4', [dim_y])
      
      y_z_xy = tf.nn.sigmoid(tf.add(tf.matmul(bn_2, w_4), b_4))
      
     #for z_x 
      w_5 = self.weight_variable('w_5', [dim_x, n_units_1])
      b_5 = self.bias_variable('b_5', [n_units_1])

      h_5 = tf.add(tf.matmul(x, w_5), b_5)
      
      # batch normalization
      beta_3 = self.beta_variable('beta_3', [n_units_1])
      gamma_3 = self.gamma_variable('gamma_3', [n_units_1])
      
      mean_3, var_3 = tf.nn.moments(h_5, [0])
      h_5 = gamma_3 * (h_5 - mean_3) / tf.sqrt(var_3 + 1e-5) + beta_3
      
      bn_3 = tf.nn.softplus(h_5)

      w_6 = self.weight_variable('w_6', [n_units_1, dim_z * 2])
      b_6 = self.bias_variable('b_6', [dim_z * 2])

      h_6 = tf.add(tf.matmul(bn_3, w_6), b_6)
      
      mu_z_x = h_6[:, : dim_z]
      logvar_z_x = h_6[:, dim_z :]
            
      z_x = mu_z_x + tf.random_normal(shape = [batch_size, dim_z], mean = 0.0, stddev = 1.0) * tf.exp(0.5 * logvar_z_x)
      
      w_7 = self.weight_variable('w_7', [dim_z, n_units_1])
      b_7 = self.bias_variable('b_7', [n_units_1])

      h_7 = tf.add(tf.matmul(z_x, w_7), b_7)

      # batch normalization
      beta_4 = self.beta_variable('beta_4', [n_units_1])
      gamma_4 = self.gamma_variable('gamma_4', [n_units_1])
      
      mean_4, var_4 = tf.nn.moments(h_7, [0])
      h_7 = gamma_4 * (h_7 - mean_4) / tf.sqrt(var_4 + 1e-5) + beta_4
      
      bn_4 = tf.nn.softplus(h_7)
      
      w_8 = self.weight_variable('w_8', [n_units_1, dim_y])
      b_8 = self.bias_variable('b_8', [dim_y])
      
      y_z_x = tf.nn.sigmoid(tf.add(tf.matmul(bn_4, w_8), b_8))
            
      return y_z_xy, y_z_x, mu_z_xy, logvar_z_xy, mu_z_x, logvar_z_x
    
  def loss(self, y_z_xy, y_z_x, y, mu_z_xy, logvar_z_xy, mu_z_x, logvar_z_x, alpha):
    likelihood_cvae =  tf.reduce_sum(y * tf.log(tf.clip_by_value(y_z_xy, 1e-10, 1.0)) + \
                                         (1 - y) * tf.log(tf.clip_by_value(1 - y_z_xy, 1e-10, 1.0)), axis = 1)
    kl_cvae = self.kl_divergence(mu_z_xy, logvar_z_xy, mu_z_x, logvar_z_x)
    elbo_cvae = tf.reduce_mean(likelihood_cvae - kl_cvae)
    loss_cvae = -elbo_cvae

    likelihood_gsnn =  tf.reduce_sum(y * tf.log(tf.clip_by_value(y_z_x, 1e-10, 1.0)) + \
                                         (1 - y) * tf.log(tf.clip_by_value(1 - y_z_x, 1e-10, 1.0)), axis = 1)
    loss_gsnn = -tf.reduce_mean(likelihood_gsnn)

    return alpha * loss_cvae + (1 - alpha) * loss_gsnn
    
  def training(self, loss, learning_rate):
    optimizer = tf.train.AdamOptimizer(learning_rate = learning_rate)
    train_step = optimizer.minimize(loss)
    return train_step
  
  def training_clipped(self, loss, learning_rate, clip_norm):
    optimizer = tf.train.AdamOptimizer(learning_rate = learning_rate)

    grads_and_vars = optimizer.compute_gradients(loss)
    clipped_grads_and_vars = [(tf.clip_by_norm(grad, clip_norm = clip_norm), \
                             var) for grad, var in grads_and_vars]
    train_step = optimizer.apply_gradients(clipped_grads_and_vars)
    
    return train_step

  def fit(self, images_train_x, images_train_y, labels_train, images_test_x, images_test_y, labels_test, \
          dim_x, dim_y, n_units_1, dim_z, alphat, \
          learning_rate, n_iter, batch_size, show_step, is_saving, model_path):
    
    tf.reset_default_graph()
    
    x = tf.placeholder(shape = [None, 14 * 28], dtype = tf.float32)
    y = tf.placeholder(shape = [None, 14 * 28], dtype = tf.float32)
    
    y_z_xy, y_z_x, mu_z_xy, logvar_z_xy, mu_z_x, logvar_z_x = self.inference(x, y, dim_x, dim_y, n_units_1, dim_z, batch_size)
    loss = self.loss(y_z_xy, y_z_x, y, mu_z_xy, logvar_z_xy, mu_z_x, logvar_z_x, alpha)
    train_step = self.training(loss, learning_rate)
    
    init = tf.global_variables_initializer()
    saver = tf.train.Saver()
    
    with tf.Session() as sess:
      
      sess.run(init)
      
      history_loss_train = []
      history_loss_test = []
      
      for i in range(n_iter):
        # Train
        rand_index = np.random.choice(len(images_train_x), size = batch_size)
        x_batch = images_train_x[rand_index]
        y_batch = images_train_y[rand_index]
        
        feed_dict = {x: x_batch, y: y_batch}

        sess.run(train_step, feed_dict = feed_dict)

        temp_loss = sess.run(loss, feed_dict = feed_dict)
        history_loss_train.append(temp_loss)

        if (i + 1) % show_step == 0:
          print ('--------------------')
          print ('Iteration: ' + str(i + 1) + '  Loss: ' + str(temp_loss))
          
        # Test
        rand_index = np.random.choice(len(images_test_x), size = batch_size)
        x_batch = images_test_x[rand_index]
        y_batch = images_test_y[rand_index]
        
        feed_dict = {x: x_batch, y: y_batch}

        temp_loss = sess.run(loss, feed_dict = feed_dict)
        history_loss_test.append(temp_loss)

      if is_saving:
        model_path = saver.save(sess, model_path)
        print ('--------------------')
        print ('done saving at ', model_path)

    fig = plt.figure(figsize = (10, 3))
    ax1 = fig.add_subplot(1, 2, 1)
    ax1.plot(range(n_iter), history_loss_train, 'b-', label = 'Train')
    ax1.plot(range(n_iter), history_loss_test, 'r--', label = 'Test')
    ax1.set_title('Loss')
    ax1.legend(loc = 'upper right')

    plt.show()    
    
  def prediction(self, x, dim_x, dim_y, n_units_1, dim_z, batch_size, model_path):
    with tf.variable_scope('enc_z_dec', reuse = True):
     #for z_x 
      w_5 = self.weight_variable('w_5', [dim_x, n_units_1])
      b_5 = self.bias_variable('b_5', [n_units_1])

      h_5 = tf.add(tf.matmul(x, w_5), b_5)
      
      # batch normalization
      beta_3 = self.beta_variable('beta_3', [n_units_1])
      gamma_3 = self.gamma_variable('gamma_3', [n_units_1])
      
      mean_3, var_3 = tf.nn.moments(h_5, [0])
      h_5 = gamma_3 * (h_5 - mean_3) / tf.sqrt(var_3 + 1e-5) + beta_3
      
      bn_3 = tf.nn.softplus(h_5)

      w_6 = self.weight_variable('w_6', [n_units_1, dim_z * 2])
      b_6 = self.bias_variable('b_6', [dim_z * 2])

      h_6 = tf.add(tf.matmul(bn_3, w_6), b_6)
      
      mu_z_x = h_6[:, : dim_z]
      logvar_z_x = h_6[:, dim_z :]
            
      z_x = mu_z_x + tf.random_normal(shape = [batch_size, dim_z], mean = 0.0, stddev = 1.0) * tf.exp(0.5 * logvar_z_x)
      
      w_7 = self.weight_variable('w_7', [dim_z, n_units_1])
      b_7 = self.bias_variable('b_7', [n_units_1])

      h_7 = tf.add(tf.matmul(z_x, w_7), b_7)

      # batch normalization
      beta_4 = self.beta_variable('beta_4', [n_units_1])
      gamma_4 = self.gamma_variable('gamma_4', [n_units_1])
      
      mean_4, var_4 = tf.nn.moments(h_7, [0])
      h_7 = gamma_4 * (h_7 - mean_4) / tf.sqrt(var_4 + 1e-5) + beta_4
      
      bn_4 = tf.nn.softplus(h_7)
      
      w_8 = self.weight_variable('w_8', [n_units_1, dim_y])
      b_8 = self.bias_variable('b_8', [dim_y])
      
      y_z_x = tf.nn.sigmoid(tf.add(tf.matmul(bn_4, w_8), b_8))

    saver = tf.train.Saver()
    
    with tf.Session() as sess:
      
      saver.restore(sess, model_path)
    
      return sess.run(y_z_x)    
```

##データ
```python
images_train_x = mnist.train.images[:, :14*28]
images_train_y = mnist.train.images[:, 14*28:]
labels_train = mnist.train.labels
images_test_x = mnist.test.images[:, :14*28]
images_test_y = mnist.test.images[:, 14*28:]
labels_test = mnist.test.labels
```

##パラメータ
dim_x = 14*28
dim_y = 14*28
n_units_1 = 64
dim_z = 2
alpha = 0.5
learning_rate = 0.01
batch_size = 100 

##アウトプット
![image.png](https://qiita-image-store.s3.amazonaws.com/0/254604/f5fea12f-9bfe-9913-7092-c400a8ef536f.png)
 
```python
index = np.random.choice(10000, batch_size)
x = images_test_x[index]
y = images_test_y[index]

predicts = cvae.prediction(x, dim_x, dim_y, n_units_1, dim_z, batch_size, model_path)

print (np.shape(x))
print (np.shape(predicts))

fig = plt.figure(figsize = (4, 1))

ax1 = fig.add_subplot(2, 3, 1)
ax1.imshow(np.reshape(x[0], [14, 28]), cmap = 'gray')
ax1.set_axis_off()
#ax1.set_title('Input')

ax2 = fig.add_subplot(2, 3, 4)
ax2.imshow(np.reshape(predicts[0], [14, 28]), cmap = 'gray')
ax2.set_axis_off()
#ax2.set_title('Predict')

ax3 = fig.add_subplot(2, 3, 2)
ax3.imshow(np.reshape(x[1], [14, 28]), cmap = 'gray')
ax3.set_axis_off()
#ax3.set_title('Input')

ax4 = fig.add_subplot(2, 3, 5)
ax4.imshow(np.reshape(predicts[1], [14, 28]), cmap = 'gray')
ax4.set_axis_off()
#ax4.set_title('Predict')

ax5 = fig.add_subplot(2, 3, 3)
ax5.imshow(np.reshape(x[2], [14, 28]), cmap = 'gray')
ax5.set_axis_off()
#ax5.set_title('Input')

ax6 = fig.add_subplot(2, 3, 6)
ax6.imshow(np.reshape(predicts[2], [14, 28]), cmap = 'gray')
ax6.set_axis_off()
#ax6.set_title('Predict')

plt.show()
```
![image.png](https://qiita-image-store.s3.amazonaws.com/0/254604/de031133-b41b-a0d6-92d3-9e6f71d882e4.png)