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: [height, width, output_channel, in_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 discriminator(input_):
with tf.variable_scope("discrimination", 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)
flatten = tf.reshape(conv2, (-1, 7*7*64))
logits = fully_contact(flatten, [7*7*64, 1], "fc")
outputs = tf.sigmoid(logits)
return outputs, logits
#------------------分割线------------------#
def generator(input_):
with tf.variable_scope("generation", reuse = tf.AUTO_REUSE) as scope1:
re_flatten = fully_contact(input_, [100, 7*7*64], "re_fc") # (?, 1)
deconv1 = tf.reshape(re_flatten, [-1, 7, 7, 64])
deconv1 = deconv2d(deconv1, [batch_size, 14, 14, 32] ,5, 2, "deconv1")
deconv2 = leakyrelu(deconv1)
deconv2 = deconv2d(deconv2, [batch_size, 28, 28, 1] ,5, 2, "deconv2")
outputs = tf.tanh(deconv2)
return outputs
#------------------分割线------------------#
input_image = tf.placeholder(tf.float32, [None, 28, 28, 1])
input_noise = tf.placeholder(tf.float32, [None, 100])
# tensor_image = tf.reshape(input_image,[-1, 28, 28, 1])
#------------------分割线------------------#
generate_image = generator(input_noise)
fake_output, fake_logits = discriminator(generate_image)
real_output, real_logits = discriminator(input_image)
#------------------分割线------------------#
G_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=fake_logits,
labels =tf.ones_like(fake_output * (1-smooth))))
D_real_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=real_logits,
labels =tf.ones_like(real_output) * (1-smooth)))
D_fake_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=fake_logits,
labels= tf.zeros_like(fake_output)))
D_loss = tf.add(D_real_loss, D_fake_loss)
#------------------分割线------------------#
train_var = tf.trainable_variables()
g_var = [var for var in train_var if var.name.startswith("generation")]
d_var = [var for var in train_var if var.name.startswith("discrimination")]
with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
train_loss_d = tf.train.AdamOptimizer(0.001, beta1=0.4).minimize(D_loss, var_list = d_var)
train_loss_g = tf.train.AdamOptimizer(0.001, beta1=0.4).minimize(G_loss, var_list = g_var)
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_noise = np.random.uniform(-1, 1, size=(batch_size, 100))
batch_image = mnist_image[0].reshape(batch_size, 28, 28, 1)
batch_image = batch_image * 2 - 1
sess.run(train_loss_d, feed_dict={input_image: batch_image, input_noise: batch_noise})
sess.run(train_loss_g, feed_dict={input_image: batch_image, input_noise: batch_noise})
if i % 10 == 0:
print(sess.run(G_loss, feed_dict={input_image: batch_image, input_noise: batch_noise}))
print(sess.run(D_loss, feed_dict={input_image: batch_image, input_noise: batch_noise}))
output_image = sess.run(generate_image, feed_dict={input_image: batch_image, input_noise: batch_noise})
save_image(output_image, 8, path, False, i)