

In [ ]:

    
#!/usr/bin/python
#
# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import

import tensorflow as tf
import numpy as np


tf.keras.backend.clear_session()

# Load data
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()


TRAINING_SIZE = len(train_images)
TEST_SIZE = len(test_images)

# Reshape from (N, 28, 28) to (N, 784)
train_images = np.reshape(train_images, (TRAINING_SIZE, 784))
test_images = np.reshape(test_images, (TEST_SIZE, 784))

# Convert the array to float32 as opposed to uint8
train_images = train_images.astype(np.float32)
test_images = test_images.astype(np.float32)

# Convert the pixel values from integers between 0 and 255 to floats between 0 and 1
train_images /= 255
test_images /=  255


NUM_DIGITS = 10

print("Before", train_labels[0]) # The format of the labels before conversion

train_labels  = tf.keras.utils.to_categorical(train_labels, NUM_DIGITS)

print("After", train_labels[0]) # The format of the labels after conversion

test_labels = tf.keras.utils.to_categorical(test_labels, NUM_DIGITS)


# Cast the labels to floats, needed later
train_labels = train_labels.astype(np.float32)
test_labels = test_labels.astype(np.float32)


# The model we implement is similar to famous LeNet5 "Gradient-Based Learning Applied to Document Recognition"
# with some simplifications:
#   (1) The first Conv2D uses default initialization.
#   (2) Use max pooling for subsampling and remove upsampling.
#   (3) Use ReLU as activation functions.

model = tf.keras.Sequential()
model.add(tf.keras.layers.Reshape((28, 28, 1), input_shape=(784,), name='reshape_1'))   # <- [None, 28, 28, 1]
model.add(tf.keras.layers.Conv2D(filters=6, kernel_size=(5, 5), padding='same', activation=tf.nn.relu, name='conv_1'))  # <- [None, 28, 28, 6]
model.add(tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2), padding='valid', name='maxpool_1'))  # <- [None, 14, 14, 6]
model.add(tf.keras.layers.Conv2D(filters=16, kernel_size=(5, 5), strides=(1, 1), padding='valid', activation=tf.nn.relu, name='conv_2'))  # <- [None, 10, 10, 16]
model.add(tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2), padding='valid', name='maxpool_2'))  # <- [None, 5, 5, 16]
model.add(tf.keras.layers.Flatten(name='flatten'))   # <- [None, 400]
model.add(tf.keras.layers.Dense(120, activation=tf.nn.relu, name='dense_1'))  # <- [None, 120]
model.add(tf.keras.layers.Dense(84, activation=tf.nn.relu, name='dense_2'))  # <- [None, 84]
model.add(tf.keras.layers.Dense(10, activation=tf.nn.softmax, name='dense_3'))  # <- [None, 10]
# Notes: Make sure each layer above has explicit names to support tfjs.

optimizer = tf.train.RMSPropOptimizer(learning_rate=0.001)

model.compile(loss='categorical_crossentropy',
              optimizer=optimizer,
              metrics=['accuracy'])

model.summary()


BATCH_SIZE=128
EPOCHS=5

model.fit(train_images, train_labels, epochs=EPOCHS, batch_size=BATCH_SIZE)

print("Starting evaluation")
loss, accuracy = model.evaluate(test_images, test_labels)
print('Test accuracy: %.2f' % (accuracy))

model.save('model.h5')