-
Notifications
You must be signed in to change notification settings - Fork 1
/
mnist-train-cnn.py
96 lines (77 loc) · 3.65 KB
/
mnist-train-cnn.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import argparse, os
import numpy as np
import tensorflow
from tensorflow.keras import backend as K
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Activation, Flatten, BatchNormalization, Conv2D, MaxPooling2D
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import multi_gpu_model
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--epochs', type=int, default=10)
parser.add_argument('--learning-rate', type=float, default=0.001)
parser.add_argument('--batch-size', type=int, default=32)
parser.add_argument('--gpu-count', type=int, default=os.environ['SM_NUM_GPUS'])
parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
parser.add_argument('--training', type=str, default=os.environ['SM_CHANNEL_TRAINING'])
parser.add_argument('--validation', type=str, default=os.environ['SM_CHANNEL_VALIDATION'])
args, _ = parser.parse_known_args()
epochs = args.epochs
lr = args.learning_rate
batch_size = args.batch_size
gpu_count = args.gpu_count
model_dir = args.model_dir
training_dir = args.training
validation_dir = args.validation
mnist_train_images = np.load(os.path.join(training_dir, 'training.npz'))['image']
mnist_train_labels = np.load(os.path.join(training_dir, 'training.npz'))['label']
mnist_test_images = np.load(os.path.join(validation_dir, 'validation.npz'))['image']
mnist_test_labels = np.load(os.path.join(validation_dir, 'validation.npz'))['label']
K.set_image_data_format('channels_last')
train_images = mnist_train_images.reshape(mnist_train_images.shape[0], 28, 28, 1)
test_images = mnist_test_images.reshape(mnist_test_images.shape[0], 28, 28, 1)
input_shape = (28, 28, 1)
train_images = train_images.astype('float32')
test_images = test_images.astype('float32')
train_images /= 255
test_images /= 255
train_labels = tensorflow.keras.utils.to_categorical(mnist_train_labels, 10)
test_labels = tensorflow.keras.utils.to_categorical(mnist_test_labels, 10)
model = Sequential()
model.add(Conv2D(32, kernel_size=(3, 3),
activation='relu',
input_shape=input_shape))
# 64 3x3 kernels
model.add(Conv2D(64, (3, 3), activation='relu'))
# Reduce by taking the max of each 2x2 block
model.add(MaxPooling2D(pool_size=(2, 2)))
# Dropout to avoid overfitting
model.add(Dropout(0.25))
# Flatten the results to one dimension for passing into our final layer
model.add(Flatten())
# A hidden layer to learn with
model.add(Dense(128, activation='relu'))
# Another dropout
model.add(Dropout(0.5))
# Final categorization from 0-9 with softmax
model.add(Dense(10, activation='softmax'))
print(model.summary())
if gpu_count > 1:
model = multi_gpu_model(model, gpus=gpu_count)
model.compile(loss=tensorflow.keras.losses.categorical_crossentropy,
optimizer=Adam(lr=lr),
metrics=['accuracy'])
model.fit(train_images, train_labels, batch_size=batch_size,
validation_data=(test_images, test_labels),
epochs=epochs,
verbose=2)
score = model.evaluate(test_images, test_labels, verbose=0)
print('Validation loss :', score[0])
print('Validation accuracy:', score[1])
# save Keras model for Tensorflow Serving
sess = K.get_session()
tensorflow.saved_model.simple_save(
sess,
os.path.join(model_dir, 'model/1'),
inputs={'inputs': model.input},
outputs={t.name: t for t in model.outputs})