|
| 1 | +# Copyright 2018 Google Inc. All Rights Reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +""" Extract from notebook for Serving Optimization on Keras """ |
| 16 | + |
| 17 | +from __future__ import print_function |
| 18 | + |
| 19 | +from datetime import datetime |
| 20 | +import os |
| 21 | +import sh |
| 22 | +import sys |
| 23 | +import tensorflow as tf |
| 24 | +from tensorflow import data |
| 25 | +from tensorflow.python.saved_model import tag_constants |
| 26 | +from tensorflow.python.tools import freeze_graph |
| 27 | +from tensorflow.python import ops |
| 28 | +from tensorflow.tools.graph_transforms import TransformGraph |
| 29 | + |
| 30 | +from inference_test import inference_test, load_mnist_keras |
| 31 | +from optimize_graph import (run_experiment, get_graph_def_from_saved_model, |
| 32 | + describe_graph, get_size, get_metagraph, get_graph_def_from_file, |
| 33 | + convert_graph_def_to_saved_model, freeze_model, optimize_graph, TRANSFORMS) |
| 34 | + |
| 35 | +NUM_CLASSES = 10 |
| 36 | +MODELS_LOCATION = 'models/mnist' |
| 37 | +MODEL_NAME = 'keras_classifier' |
| 38 | + |
| 39 | + |
| 40 | +def keras_model_fn(params): |
| 41 | + |
| 42 | + inputs = tf.keras.layers.Input(shape=(28, 28), name='input_image') |
| 43 | + input_layer = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name='reshape')(inputs) |
| 44 | + |
| 45 | + # convolutional layers |
| 46 | + conv_inputs = input_layer |
| 47 | + for i in range(params.num_conv_layers): |
| 48 | + filters = params.init_filters * (2**i) |
| 49 | + conv = tf.keras.layers.Conv2D(kernel_size=3, filters=filters, strides=1, padding='SAME', activation='relu')(conv_inputs) |
| 50 | + max_pool = tf.keras.layers.MaxPool2D(pool_size=2, strides=2, padding='SAME')(conv) |
| 51 | + batch_norm = tf.keras.layers.BatchNormalization()(max_pool) |
| 52 | + conv_inputs = batch_norm |
| 53 | + |
| 54 | + flatten = tf.keras.layers.Flatten(name='flatten')(conv_inputs) |
| 55 | + |
| 56 | + # fully-connected layers |
| 57 | + dense_inputs = flatten |
| 58 | + for i in range(len(params.hidden_units)): |
| 59 | + dense = tf.keras.layers.Dense(units=params.hidden_units[i], activation='relu')(dense_inputs) |
| 60 | + dropout = tf.keras.layers.Dropout(params.dropout)(dense) |
| 61 | + dense_inputs = dropout |
| 62 | + |
| 63 | + # softmax classifier |
| 64 | + logits = tf.keras.layers.Dense(units=NUM_CLASSES, name='logits')(dense_inputs) |
| 65 | + softmax = tf.keras.layers.Activation('softmax', name='softmax')(logits) |
| 66 | + |
| 67 | + # keras model |
| 68 | + model = tf.keras.models.Model(inputs, softmax) |
| 69 | + return model |
| 70 | + |
| 71 | + |
| 72 | +def create_estimator_keras(params, run_config): |
| 73 | + |
| 74 | + keras_model = keras_model_fn(params) |
| 75 | + print(keras_model.summary()) |
| 76 | + |
| 77 | + optimizer = tf.keras.optimizers.Adam(lr=params.learning_rate) |
| 78 | + keras_model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy']) |
| 79 | + mnist_classifier = tf.keras.estimator.model_to_estimator( |
| 80 | + keras_model=keras_model, |
| 81 | + config=run_config |
| 82 | + ) |
| 83 | + |
| 84 | + return mnist_classifier |
| 85 | + |
| 86 | + |
| 87 | +#### Train and Export Model |
| 88 | + |
| 89 | +def train_and_export_model(train_data, train_labels): |
| 90 | + model_dir = os.path.join(MODELS_LOCATION, MODEL_NAME) |
| 91 | + |
| 92 | + hparams = tf.contrib.training.HParams( |
| 93 | + batch_size=100, |
| 94 | + hidden_units=[512, 512], |
| 95 | + num_conv_layers=3, |
| 96 | + init_filters=64, |
| 97 | + dropout=0.2, |
| 98 | + max_training_steps=50, |
| 99 | + eval_throttle_secs=10, |
| 100 | + learning_rate=1e-3, |
| 101 | + debug=True |
| 102 | + ) |
| 103 | + |
| 104 | + run_config = tf.estimator.RunConfig( |
| 105 | + tf_random_seed=19830610, |
| 106 | + save_checkpoints_steps=1000, |
| 107 | + keep_checkpoint_max=3, |
| 108 | + model_dir=model_dir |
| 109 | + ) |
| 110 | + |
| 111 | + if tf.gfile.Exists(model_dir): |
| 112 | + print('Removing previous artifacts...') |
| 113 | + tf.gfile.DeleteRecursively(model_dir) |
| 114 | + |
| 115 | + os.makedirs(model_dir) |
| 116 | + |
| 117 | + estimator = run_experiment(hparams, train_data, train_labels, run_config, create_estimator_keras) |
| 118 | + |
| 119 | + def make_serving_input_receiver_fn(): |
| 120 | + inputs = {'input_image': tf.placeholder( |
| 121 | + shape=[None,28,28], dtype=tf.float32, name='serving_input_image')} |
| 122 | + return tf.estimator.export.build_raw_serving_input_receiver_fn(inputs) |
| 123 | + |
| 124 | + export_dir = os.path.join(model_dir, 'export') |
| 125 | + |
| 126 | + if tf.gfile.Exists(export_dir): |
| 127 | + tf.gfile.DeleteRecursively(export_dir) |
| 128 | + |
| 129 | + estimator.export_savedmodel( |
| 130 | + export_dir_base=export_dir, |
| 131 | + serving_input_receiver_fn=make_serving_input_receiver_fn() |
| 132 | + ) |
| 133 | + |
| 134 | + return export_dir |
| 135 | + |
| 136 | + |
| 137 | +def setup_model(): |
| 138 | + train_data, train_labels, eval_data, eval_labels = load_mnist_keras() |
| 139 | + export_dir = train_and_export_model(train_data, train_labels) |
| 140 | + return export_dir, eval_data |
| 141 | + |
| 142 | + |
| 143 | +NUM_TRIALS = 10 |
| 144 | + |
| 145 | +def main(args): |
| 146 | + if len(args) > 1 and args[1] == '--inference': |
| 147 | + export_dir = args[2] |
| 148 | + _, _, eval_data, _ = load_mnist_keras() |
| 149 | + |
| 150 | + total_load_time = 0.0 |
| 151 | + total_serve_time = 0.0 |
| 152 | + saved_model_dir = os.path.join( |
| 153 | + export_dir, [f for f in os.listdir(export_dir) if f.isdigit()][0]) |
| 154 | + for i in range(0, NUM_TRIALS): |
| 155 | + load_time, serving_time = inference_test(saved_model_dir, eval_data, repeat=10000) |
| 156 | + total_load_time += load_time |
| 157 | + total_serve_time += serving_time |
| 158 | + |
| 159 | + print("****************************************") |
| 160 | + print("*** Load time on original model: {:.2f}".format(total_load_time / NUM_TRIALS)) |
| 161 | + print("*** Serve time on original model: {:.2f}".format(total_serve_time / NUM_TRIALS)) |
| 162 | + print("****************************************") |
| 163 | + |
| 164 | + total_load_time = 0.0 |
| 165 | + total_serve_time = 0.0 |
| 166 | + optimized_export_dir = os.path.join(export_dir, 'optimized') |
| 167 | + for i in range(0, NUM_TRIALS): |
| 168 | + load_time, serving_time = inference_test(optimized_export_dir, eval_data, |
| 169 | + signature='serving_default', |
| 170 | + repeat=10000) |
| 171 | + total_load_time += load_time |
| 172 | + total_serve_time += serving_time |
| 173 | + print("****************************************") |
| 174 | + print("*** Load time on optimized model: {:.2f}".format(total_load_time / NUM_TRIALS)) |
| 175 | + print("*** Serve time on optimized model: {:.2f}".format(total_serve_time / NUM_TRIALS)) |
| 176 | + print("****************************************") |
| 177 | + |
| 178 | + else: |
| 179 | + # generate and output original model |
| 180 | + export_dir, eval_data = setup_model() |
| 181 | + saved_model_dir = os.path.join(export_dir, os.listdir(export_dir)[-1]) |
| 182 | + describe_graph(get_graph_def_from_saved_model(saved_model_dir)) |
| 183 | + get_size(saved_model_dir, 'saved_model.pb') |
| 184 | + get_metagraph(saved_model_dir) |
| 185 | + |
| 186 | + # freeze model and describe it |
| 187 | + freeze_model(saved_model_dir, 'softmax/Softmax', 'frozen_model.pb') |
| 188 | + frozen_filepath = os.path.join(saved_model_dir, 'frozen_model.pb') |
| 189 | + describe_graph(get_graph_def_from_file(frozen_filepath)) |
| 190 | + get_size(saved_model_dir, 'frozen_model.pb', include_vars=False) |
| 191 | + |
| 192 | + # optimize model and describe it |
| 193 | + optimize_graph(saved_model_dir, 'frozen_model.pb', TRANSFORMS, 'softmax/Softmax') |
| 194 | + optimized_filepath = os.path.join(saved_model_dir, 'optimized_model.pb') |
| 195 | + describe_graph(get_graph_def_from_file(optimized_filepath)) |
| 196 | + get_size(saved_model_dir, 'optimized_model.pb', include_vars=False) |
| 197 | + |
| 198 | + # convert to saved model and output metagraph again |
| 199 | + optimized_export_dir = os.path.join(export_dir, 'optimized') |
| 200 | + convert_graph_def_to_saved_model(optimized_export_dir, optimized_filepath, |
| 201 | + 'softmax', 'softmax/Softmax:0') |
| 202 | + get_size(optimized_export_dir, 'saved_model.pb') |
| 203 | + get_metagraph(optimized_export_dir) |
| 204 | + |
| 205 | + |
| 206 | +if __name__ == '__main__': |
| 207 | + main(sys.argv) |
0 commit comments