|
| 1 | +from __future__ import print_function |
| 2 | +import shutil |
| 3 | +import os.path |
| 4 | +import tensorflow as tf |
| 5 | +from tensorflow.examples.tutorials.mnist import input_data |
| 6 | + |
| 7 | +EXPORT_DIR = './model' |
| 8 | + |
| 9 | +if os.path.exists(EXPORT_DIR): |
| 10 | + shutil.rmtree(EXPORT_DIR) |
| 11 | + |
| 12 | +mnist = input_data.read_data_sets("MNIST_data/", one_hot=True) |
| 13 | + |
| 14 | +# Parameters |
| 15 | +learning_rate = 0.001 |
| 16 | +training_iters = 200000 |
| 17 | +batch_size = 128 |
| 18 | +display_step = 10 |
| 19 | + |
| 20 | +# Network Parameters |
| 21 | +n_input = 784 # MNIST data input (img shape: 28*28) |
| 22 | +n_classes = 10 # MNIST total classes (0-9 digits) |
| 23 | +dropout = 0.75 # Dropout, probability to keep units |
| 24 | + |
| 25 | +# tf Graph input |
| 26 | +x = tf.placeholder(tf.float32, [None, n_input]) |
| 27 | +y = tf.placeholder(tf.float32, [None, n_classes]) |
| 28 | +keep_prob = tf.placeholder(tf.float32) # dropout (keep probability) |
| 29 | + |
| 30 | + |
| 31 | +# Create some wrappers for simplicity |
| 32 | +def conv2d(x, W, b, strides=1): |
| 33 | + # Conv2D wrapper, with bias and relu activation |
| 34 | + x = tf.nn.conv2d(x, W, strides=[1, strides, strides, 1], padding='SAME') |
| 35 | + x = tf.nn.bias_add(x, b) |
| 36 | + return tf.nn.relu(x) |
| 37 | + |
| 38 | + |
| 39 | +def maxpool2d(x, k=2): |
| 40 | + # MaxPool2D wrapper |
| 41 | + return tf.nn.max_pool(x, ksize=[1, k, k, 1], strides=[1, k, k, 1], |
| 42 | + padding='SAME') |
| 43 | + |
| 44 | + |
| 45 | +# Create Model |
| 46 | +def conv_net(x, weights, biases, dropout): |
| 47 | + # Reshape input picture |
| 48 | + x = tf.reshape(x, shape=[-1, 28, 28, 1]) |
| 49 | + |
| 50 | + # Convolution Layer |
| 51 | + conv1 = conv2d(x, weights['wc1'], biases['bc1']) |
| 52 | + # Max Pooling (down-sampling) |
| 53 | + conv1 = maxpool2d(conv1, k=2) |
| 54 | + |
| 55 | + # Convolution Layer |
| 56 | + conv2 = conv2d(conv1, weights['wc2'], biases['bc2']) |
| 57 | + # Max Pooling (down-sampling) |
| 58 | + conv2 = maxpool2d(conv2, k=2) |
| 59 | + |
| 60 | + # Fully connected layer |
| 61 | + # Reshape conv2 output to fit fully connected layer input |
| 62 | + fc1 = tf.reshape(conv2, [-1, weights['wd1'].get_shape().as_list()[0]]) |
| 63 | + fc1 = tf.add(tf.matmul(fc1, weights['wd1']), biases['bd1']) |
| 64 | + fc1 = tf.nn.relu(fc1) |
| 65 | + # Apply Dropout |
| 66 | + fc1 = tf.nn.dropout(fc1, dropout) |
| 67 | + |
| 68 | + # Output, class prediction |
| 69 | + out = tf.add(tf.matmul(fc1, weights['out']), biases['out']) |
| 70 | + return out |
| 71 | + |
| 72 | + |
| 73 | +# Store layers weight & bias |
| 74 | +weights = { |
| 75 | + # 5x5 conv, 1 input, 32 outputs |
| 76 | + 'wc1': tf.Variable(tf.random_normal([5, 5, 1, 32])), |
| 77 | + # 5x5 conv, 32 inputs, 64 outputs |
| 78 | + 'wc2': tf.Variable(tf.random_normal([5, 5, 32, 64])), |
| 79 | + # fully connected, 7*7*64 inputs, 1024 outputs |
| 80 | + 'wd1': tf.Variable(tf.random_normal([7 * 7 * 64, 1024])), |
| 81 | + # 1024 inputs, 10 outputs (class prediction) |
| 82 | + 'out': tf.Variable(tf.random_normal([1024, n_classes])) |
| 83 | +} |
| 84 | + |
| 85 | +biases = { |
| 86 | + 'bc1': tf.Variable(tf.random_normal([32])), |
| 87 | + 'bc2': tf.Variable(tf.random_normal([64])), |
| 88 | + 'bd1': tf.Variable(tf.random_normal([1024])), |
| 89 | + 'out': tf.Variable(tf.random_normal([n_classes])) |
| 90 | +} |
| 91 | + |
| 92 | +# Construct model |
| 93 | +pred = conv_net(x, weights, biases, keep_prob) |
| 94 | + |
| 95 | +# Define loss and optimizer |
| 96 | +cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(pred, y)) |
| 97 | +optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost) |
| 98 | + |
| 99 | +# Evaluate model |
| 100 | +correct_pred = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1)) |
| 101 | +accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32)) |
| 102 | + |
| 103 | +# Initializing the variables |
| 104 | +init = tf.initialize_all_variables() |
| 105 | + |
| 106 | +# Launch the graph |
| 107 | +with tf.Session() as sess: |
| 108 | + sess.run(init) |
| 109 | + step = 1 |
| 110 | + # Keep training until reach max iterations |
| 111 | + while step * batch_size < training_iters: |
| 112 | + batch_x, batch_y = mnist.train.next_batch(batch_size) |
| 113 | + # Run optimization op (backprop) |
| 114 | + sess.run(optimizer, feed_dict={x: batch_x, y: batch_y, |
| 115 | + keep_prob: dropout}) |
| 116 | + if step % display_step == 0: |
| 117 | + # Calculate batch loss and accuracy |
| 118 | + loss, acc = sess.run([cost, accuracy], feed_dict={x: batch_x, |
| 119 | + y: batch_y, |
| 120 | + keep_prob: 1.}) |
| 121 | + print("Iter " + str(step * batch_size) + ", Minibatch Loss= " + \ |
| 122 | + "{:.6f}".format(loss) + ", Training Accuracy= " + \ |
| 123 | + "{:.5f}".format(acc)) |
| 124 | + step += 1 |
| 125 | + print("Optimization Finished!") |
| 126 | + |
| 127 | + # Calculate accuracy for 256 mnist test images |
| 128 | + print("Testing Accuracy:", \ |
| 129 | + sess.run(accuracy, feed_dict={x: mnist.test.images[:256], |
| 130 | + y: mnist.test.labels[:256], |
| 131 | + keep_prob: 1.})) |
| 132 | + WC1 = weights['wc1'].eval(sess) |
| 133 | + BC1 = biases['bc1'].eval(sess) |
| 134 | + WC2 = weights['wc2'].eval(sess) |
| 135 | + BC2 = biases['bc2'].eval(sess) |
| 136 | + WD1 = weights['wd1'].eval(sess) |
| 137 | + BD1 = biases['bd1'].eval(sess) |
| 138 | + W_OUT = weights['out'].eval(sess) |
| 139 | + B_OUT = biases['out'].eval(sess) |
| 140 | + |
| 141 | +# Create new graph for exporting |
| 142 | +g = tf.Graph() |
| 143 | +with g.as_default(): |
| 144 | + x_2 = tf.placeholder("float", shape=[None, 784], name="input") |
| 145 | + |
| 146 | + WC1 = tf.constant(WC1, name="WC1") |
| 147 | + BC1 = tf.constant(BC1, name="BC1") |
| 148 | + x_image = tf.reshape(x_2, [-1, 28, 28, 1]) |
| 149 | + CONV1 = conv2d(x_image, WC1, BC1) |
| 150 | + MAXPOOL1 = maxpool2d(CONV1, k=2) |
| 151 | + |
| 152 | + WC2 = tf.constant(WC2, name="WC2") |
| 153 | + BC2 = tf.constant(BC2, name="BC2") |
| 154 | + CONV2 = conv2d(MAXPOOL1, WC2, BC2) |
| 155 | + MAXPOOL2 = maxpool2d(CONV2, k=2) |
| 156 | + |
| 157 | + WD1 = tf.constant(WD1, name="WD1") |
| 158 | + BD1 = tf.constant(BD1, name="BD1") |
| 159 | + |
| 160 | + FC1 = tf.reshape(MAXPOOL2, [-1, WD1.get_shape().as_list()[0]]) |
| 161 | + FC1 = tf.add(tf.matmul(FC1, WD1), BD1) |
| 162 | + FC1 = tf.nn.relu(FC1) |
| 163 | + |
| 164 | + W_OUT = tf.constant(W_OUT, name="W_OUT") |
| 165 | + B_OUT = tf.constant(B_OUT, name="B_OUT") |
| 166 | + |
| 167 | + # skipped dropout for exported graph. |
| 168 | + |
| 169 | + OUTPUT = tf.nn.softmax(tf.matmul(FC1, W_OUT) + B_OUT, name="output") |
| 170 | + |
| 171 | + sess = tf.Session() |
| 172 | + init = tf.initialize_all_variables() |
| 173 | + sess.run(init) |
| 174 | + |
| 175 | + graph_def = g.as_graph_def() |
| 176 | + tf.train.write_graph(graph_def, EXPORT_DIR, 'mnist_model_graph.pb', as_text=False) |
| 177 | + |
| 178 | + # Test trained model |
| 179 | + y_train = tf.placeholder("float", [None, 10]) |
| 180 | + correct_prediction = tf.equal(tf.argmax(OUTPUT, 1), tf.argmax(y_train, 1)) |
| 181 | + accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float")) |
| 182 | + |
| 183 | + print("check accuracy %g" % accuracy.eval( |
| 184 | + {x_2: mnist.test.images, y_train: mnist.test.labels}, sess)) |
0 commit comments