Skip to content

Instantly share code, notes, and snippets.

@hgaiser
Last active July 13, 2017 07:29
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save hgaiser/9d38ac49424f400c161c8086ff864d08 to your computer and use it in GitHub Desktop.
MNIST example with Keras model and TF ops
#!/usr/bin/env python
import tensorflow as tf
from tensorflow.contrib.keras.api import keras
from tensorflow.contrib.keras.api.keras.models import Model, load_model
from tensorflow.contrib.keras.api.keras.layers import Input, Dense, Dropout, Flatten, Conv2D, MaxPooling2D, Activation, Lambda
from tensorflow.contrib.keras.api.keras.datasets import mnist
from tensorflow.contrib.keras.api.keras.utils import to_categorical
import numpy as np
import matplotlib.pyplot as plt
import argparse
def get_session():
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
return tf.Session(config=config)
def weight_variable(shape):
initial = tf.truncated_normal(shape, stddev=0.1)
return tf.Variable(initial)
def bias_variable(shape):
initial = tf.constant(0.1, shape=shape)
return tf.Variable(initial)
def conv2d(x, W):
return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')
def max_pool_2x2(x):
return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
def create_model(num_classes):
inputs = Input(shape=(28, 28, 1))
W_conv1 = weight_variable([3, 3, 1, 32])
b_conv1 = bias_variable([32])
W_conv2 = weight_variable([3, 3, 32, 64])
b_conv2 = bias_variable([64])
W_fc1 = weight_variable([14 * 14 * 64, 128])
b_fc1 = bias_variable([128])
W_fc2 = weight_variable([128, 10])
b_fc2 = bias_variable([10])
x = Lambda(lambda x: tf.nn.relu(conv2d(x, W_conv1) + b_conv1))(inputs)
x = Lambda(lambda x: max_pool_2x2(tf.nn.relu(conv2d(x, W_conv2) + b_conv2)))(x)
x = Dropout(0.25)(x)
x = Flatten()(x)
x = Lambda(lambda x: tf.nn.relu(tf.matmul(x, W_fc1) + b_fc1))(x)
x = Dropout(0.5)(x)
x = Lambda(lambda x: tf.matmul(x, W_fc2) + b_fc2)(x)
predictions = Activation('softmax')(x)
model = Model(inputs=inputs, outputs=predictions)
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
print(model.summary())
return model
def parse_args():
parser = argparse.ArgumentParser(description='MNIST Training.')
parser.add_argument('--model', help='Model to use to resume training.')
parser.add_argument('--no-train', help='Disables training, only evaluate.', dest='no_train', action='store_true')
parser.add_argument('--num-classes', help='Number of classes.', default=10)
parser.add_argument('--seed', help='Random seed to use.', default=1)
parser.add_argument('--batch-size', help='Size of the batch.', default=128)
parser.add_argument('--epochs', help='Number of epochs to run.', default=4)
parser.add_argument('--target', help='Target file for the resulting model.', default='mnist.h5')
parser.set_defaults(no_train=False)
return parser.parse_args()
if __name__=='__main__':
args = parse_args()
np.random.seed(args.seed)
keras.backend.set_session(get_session())
# load and preprocess data
(X_train, y_train), (X_test, y_test) = mnist.load_data()
X_train = X_train.reshape(-1, 28, 28, 1).astype(np.float32) / 255
X_test = X_test.reshape(-1, 28, 28, 1).astype(np.float32) / 255
y_train = to_categorical(y_train, args.num_classes)
y_test = to_categorical(y_test, args.num_classes)
if args.model:
print('Loading model from {}'.format(args.model))
model = load_model(args.model)
else:
# create model
model = create_model(args.num_classes)
if not args.no_train:
# fit model
model.fit(
X_train,
y_train,
batch_size=args.batch_size,
epochs=args.epochs,
shuffle=True,
verbose=1,
validation_data=(X_test, y_test)
)
# save model
model.save(args.target)
# evaluate model
score = model.evaluate(X_test, y_test, verbose=0)
print('Test loss: {}'.format(score[0]))
print('Test accuracy: {}'.format(score[1]))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment