tensorflow 卷积、反卷积形式的去噪自编码器

1. #! -*- coding: utf-8 -*-
2. ## by Colie (lijixiang)
3. import tensorflow as tf
4. from tensorflow.examples.tutorials.mnist import input_data
5. import numpy as np
6. from PIL import Image
7. train_epochs = 35  ## int(1e5+1)
8. INPUT_HEIGHT = 28
9. INPUT_WIDTH = 28
10. batch_size = 256
11. noise_factor = 0.5  ## (0~1)
12. ## 原始输入是28×28*3
13. input_x = tf.placeholder(tf.float32, [None, INPUT_HEIGHT * INPUT_WIDTH], name='input_with_noise')
14. input_matrix = tf.reshape(input_x, shape=[-1, INPUT_HEIGHT, INPUT_WIDTH, 1])
15. input_raw = tf.placeholder(tf.float32, shape=[None, INPUT_HEIGHT * INPUT_WIDTH], name='input_without_noise')
16. ## 1 conv layer
17. ## 输入28*28*3
18. ## 经过卷积、激活、池化，输出14*14*64
19. weight_1 = tf.Variable(tf.truncated_normal(shape=[3, 3, 1, 64], stddev=0.1, name = 'weight_1'))
20. bias_1 = tf.Variable(tf.constant(0.0, shape=[64], name='bias_1'))
21. conv1 = tf.nn.conv2d(input=input_matrix, filter=weight_1, strides=[1, 1, 1, 1], padding='SAME')
22. conv1 = tf.nn.bias_add(conv1, bias_1, name='conv_1')
23. acti1 = tf.nn.relu(conv1, name='acti_1')
24. pool1 = tf.nn.max_pool(value=acti1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name='max_pool_1')
25. ## 2 conv layer
26. ## 输入14*14*64
27. ## 经过卷积、激活、池化，输出7×7×64
28. weight_2 = tf.Variable(tf.truncated_normal(shape=[3, 3, 64, 64], stddev=0.1, name='weight_2'))
29. bias_2 = tf.Variable(tf.constant(0.0, shape=[64], name='bias_2'))
30. conv2 = tf.nn.conv2d(input=pool1, filter=weight_2, strides=[1, 1, 1, 1], padding='SAME')
31. conv2 = tf.nn.bias_add(conv2, bias_2, name='conv_2')
32. acti2 = tf.nn.relu(conv2, name='acti_2')
33. pool2 = tf.nn.max_pool(value=acti2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name='max_pool_2')
34. ## 3 conv layer
35. ## 输入7*7*64
36. ## 经过卷积、激活、池化，输出4×4×32
37. ## 原始输入是28*28*3=2352，转化为4*4*32=512，大量噪声会在网络中过滤掉
38. weight_3 = tf.Variable(tf.truncated_normal(shape=[3, 3, 64, 32], stddev=0.1, name='weight_3'))
39. bias_3 = tf.Variable(tf.constant(0.0, shape=[32]))
40. conv3 = tf.nn.conv2d(input=pool2, filter=weight_3, strides=[1, 1, 1, 1], padding='SAME')
41. conv3 = tf.nn.bias_add(conv3, bias_3)
42. acti3 = tf.nn.relu(conv3, name='acti_3')
43. pool3 = tf.nn.max_pool(value=acti3, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name='max_pool_3')
44. ## 1 deconv layer
45. ## 输入4*4*32
46. ## 经过反卷积，输出7*7*32
47. deconv_weight_1 = tf.Variable(tf.truncated_normal(shape=[3, 3, 32, 32], stddev=0.1), name='deconv_weight_1')
48. deconv1 = tf.nn.conv2d_transpose(value=pool3, filter=deconv_weight_1, output_shape=[batch_size, 7, 7, 32], strides=[1, 2, 2, 1], padding='SAME', name='deconv_1')
49. ## 2 deconv layer
50. ## 输入7*7*32
51. ## 经过反卷积，输出14*14*64
52. deconv_weight_2 = tf.Variable(tf.truncated_normal(shape=[3, 3, 64, 32], stddev=0.1), name='deconv_weight_2')
53. deconv2 = tf.nn.conv2d_transpose(value=deconv1, filter=deconv_weight_2, output_shape=[batch_size, 14, 14, 64], strides=[1, 2, 2, 1], padding='SAME', name='deconv_2')
54. ## 3 deconv layer
55. ## 输入14*14*64
56. ## 经过反卷积，输出28*28*64
57. deconv_weight_3 = tf.Variable(tf.truncated_normal(shape=[3, 3, 64, 64], stddev=0.1, name='deconv_weight_3'))
58. deconv3 = tf.nn.conv2d_transpose(value=deconv2, filter=deconv_weight_3, output_shape=[batch_size, 28, 28, 64], strides=[1, 2, 2, 1], padding='SAME', name='deconv_3')
59. ## conv layer
60. ## 输入28*28*64
61. ## 经过卷积，输出为28*28*1
62. weight_final = tf.Variable(tf.truncated_normal(shape=[3, 3, 64, 1], stddev=0.1, name = 'weight_final'))
63. bias_final = tf.Variable(tf.constant(0.0, shape=[1], name='bias_final'))
64. conv_final = tf.nn.conv2d(input=deconv3, filter=weight_final, strides=[1, 1, 1, 1], padding='SAME')
65. conv_final = tf.nn.bias_add(conv_final, bias_final, name='conv_final')
66. ## output
67. ## 输入28*28*1
68. ## reshape为28*28
69. output = tf.reshape(conv_final, shape=[-1, INPUT_HEIGHT * INPUT_WIDTH])
70. ## loss and optimizer
71. loss = tf.reduce_mean(tf.pow(tf.subtract(output, input_raw), 2.0))
72. optimizer = tf.train.AdamOptimizer(0.01).minimize(loss)
73. with tf.Session() as sess:
74. mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
75. n_samples = int(mnist.train.num_examples)
76. print('train samples: %d' % n_samples)
77. print('batch size: %d' % batch_size)
78. total_batch = int(n_samples / batch_size)
79. print('total batchs: %d' % total_batch)
80. init = tf.global_variables_initializer()
81. sess.run(init)
82. for epoch in range(train_epochs):
83. for batch_index in range(total_batch):
84. batch_x, _ = mnist.train.next_batch(batch_size)
85. noise_x = batch_x + noise_factor * np.random.randn(*batch_x.shape)
86. noise_x = np.clip(noise_x, 0., 1.)
87. _, train_loss = sess.run([optimizer, loss], feed_dict={input_x: noise_x, input_raw: batch_x})
88. print('epoch: %04d\tbatch: %04d\ttrain loss: %.9f' % (epoch + 1, batch_index + 1, train_loss))
89. ## 训练结束后，用测试集测试，并保存加噪图像、去噪图像
90. n_test_samples = int(mnist.test.num_examples)
91. test_total_batch = int(n_test_samples / batch_size)
92. for i in range(test_total_batch):
93. batch_test_x, _ = mnist.test.next_batch(batch_size)
94. noise_test_x = batch_test_x + noise_factor * np.random.randn(*batch_test_x.shape)
95. noise_test_x = np.clip(noise_test_x, 0., 1.)
96. test_loss, pred_result = sess.run([loss, conv_final], feed_dict={input_x: noise_test_x, input_raw: batch_test_x})
97. print('test batch index: %d\ttest loss: %.9f' % (i + 1, test_loss))
98. for index in range(batch_size):
99. array = np.reshape(pred_result[index], newshape=[INPUT_HEIGHT, INPUT_WIDTH])
100. array = array * 255
101. image = Image.fromarray(array)
102. if image.mode != 'L':
103. image = image.convert('L')
104. image.save('./pred/' + str(i * batch_size + index) + '.png')
105. array_raw = np.reshape(noise_test_x[index], newshape=[INPUT_HEIGHT, INPUT_WIDTH])
106. array_raw = array_raw * 255
107. image_raw = Image.fromarray(array_raw)
108. if image_raw.mode != 'L':
109. image_raw = image_raw.convert('L')
110. image_raw.save('./pred/' + str(i * batch_size + index) + '_raw.png')
111. #break</span>