tensorflow AutoEncoder

import tensorflow as tf

import numpy as np

import tensorflow.examples.tutorials.mnist.input_data as input_data

import math

import scipy

mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
#------------------分割线------------------#

batch_size=64

smooth=0.1

path = '/Users/changxingya/Documents/test_image'

#------------------分割线------------------#
#use the msra to initialize the parameters
'''
input_: Tensor of input, format NHWC

shape: [filter_height, filter_width, in_channel, out_channel]

k_step: The size of kernel

name: Name of model

'''
def conv2d(input_, shape, k_step, name):

  with tf.variable_scope(name):

    msra_num = 1.0

    fan_in = k_step * k_step * int(input_.get_shape()[-1])

    stddev = msra_num * (math.sqrt(2. / float(fan_in)))

    w = tf.get_variable('w', shape, initializer= tf.truncated_normal_initializer(stddev=stddev))

    b = tf.get_variable('b', [shape[-1]], initializer=tf.constant_initializer(value=0.0))

    conv = tf.nn.conv2d(input_, w, strides=[1,k_step,k_step,1], padding = 'SAME') + b

    return conv

'''
  input_: Tensor of input, format NHWC

  output: [batch, out_height, out_width, out_channel]

  k_step: The size of kernel is k_step*k_step

  d_step: we can define the generated image size, if you define the d_step=2,
          we can get the double size of generated image

'''

def deconv2d(input_, out_shape, k_step, d_step, name):

  with tf.variable_scope(name):

    msra_num = 1.0

    fan_in = k_step * k_step * int(input_.get_shape()[-1])

    stddev = msra_num * (math.sqrt(2. / float(fan_in) * float(d_step) * float(d_step)))

    w = tf.get_variable('w', [k_step, k_step, out_shape[-1], input_.get_shape()[-1]])

    deconv = tf.nn.conv2d_transpose(input_, w, output_shape=out_shape, strides=[1, d_step, d_step, 1])

  return deconv

'''

input_: Tensor of input, format NHWC

shape: [input_channel, out_channel]

'''

def fully_contact(input_, shape, name):

  with tf.variable_scope(name):

    msra_num = 1.0

    fan_in = int(input_.get_shape()[-1])

    stddev = msra_num * (math.sqrt(2./float(fan_in)))

    w = tf.get_variable('w', shape, initializer= tf.truncated_normal_initializer(stddev=stddev))

    b = tf.get_variable('b', shape[-1], initializer= tf.constant_initializer(value=0.0))

    fc = tf.matmul(input_, w) + b

    return fc

'''

define the function of leakrelu

'''

def leakyrelu(x, leak=0.2):

  k1 = (1 + leak)*0.5

  k2 = (1 - leak)*0.5

  return k1 * x + k2 * tf.abs(x)

'''

Restore pixel [-1, 1] to [0, 255]

'''

def rescale_image(image):

  convert_image = (image / 1.5 + 0.5) * 255

  return convert_image

'''

input: The tensor of input, format NHWC

size: recevie the number of images, such as size=8, we can get the 64 images, simultaneously

image_path: the path to store image

colorL: Ture is color image, Flase is gray image

iter: record continous storage images

'''

def save_image(input_, size, image_path, color, iter):

  h, w = input_.shape[1],input_.shape[2]

  convert_input = input_.reshape(batch_size, h, w)

  if color is True:

    image = np.zeros((h * size, w * size, 3))

  else:

    image = np.zeros((h * size, w * size))

  for index, img in enumerate(convert_input):

    i = index % size

    j = math.floor(index / size)

    if color is True:

      image[h*j:h*j+h, i*w:i*w+w,:] = img

    else:

      image[h*j:h*j+h, i*w:i*w+w] = img

  scipy.misc.toimage(rescale_image(image),cmin=0, cmax=255).save(image_path+'/tr_gt_%s.png' % (iter))



#------------------分割线------------------#

def AutoEncoder(input_):

  with tf.variable_scope("Auotencoder", reuse = tf.AUTO_REUSE) as scope0:


      conv1 = conv2d(input_, [5, 5, 1, 32], 2, "conv1")

      conv1 = leakyrelu(conv1)

      conv2 = conv2d(conv1, [5, 5, 32, 64], 2, "conv2")

      conv2 = leakyrelu(conv2)

      deconv1 = deconv2d(conv2, [batch_size, 14, 14, 32], 5, 2, "deconv1")

      deconv1 = leakyrelu(deconv1)

      deconv2 = deconv2d(deconv1, [batch_size, 28, 28, 1], 5, 2, "deconv2")

      output = tf.tanh(deconv2)

      return output

#------------------分割线------------------#

input_image = tf.placeholder(tf.float32, [None, 28, 28, 1])

#------------------分割线------------------#

generate_image = AutoEncoder(input_image)

#------------------分割线------------------#

Auto_loss = tf.reduce_mean(tf.reduce_sum(tf.pow(tf.subtract(generate_image, input_image), 2), 3))

#------------------分割线------------------#

train_var = tf.trainable_variables()

with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):

  train_loss = tf.train.AdamOptimizer(0.001, beta1=0.9).minimize(Auto_loss)

init = tf.global_variables_initializer()

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.6)

#------------------分割线------------------#

with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:

  sess.run(init)

  for i in range(1000):

    mnist_image= mnist.train.next_batch(batch_size)

    batch_image = mnist_image[0].reshape(batch_size, 28, 28, 1)

    batch_image = batch_image * 2 - 1

    sess.run(train_loss, feed_dict={input_image: batch_image})

    if i % 10 == 0:

      print(sess.run(Auto_loss, feed_dict={input_image: batch_image}))

      output_image = sess.run(generate_image, feed_dict={input_image: batch_image})

      save_image(output_image, 8, path, False, i)