Repository: leimao/Frozen_Graph_TensorFlow Branch: master Commit: 8767d1ee8aa5 Files: 15 Total size: 44.3 KB Directory structure: gitextract_4vgmz1r4/ ├── LICENSE.md ├── README.md ├── TensorFlow_v1/ │ ├── .gitignore │ ├── README.md │ ├── cifar.py │ ├── cnn.py │ ├── inspect_signature.py │ ├── main.py │ ├── test_pb.py │ └── utils.py └── TensorFlow_v2/ ├── .gitignore ├── README.md ├── example_1.py ├── example_2.py └── utils.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: LICENSE.md ================================================ The MIT License (MIT) Copyright (c) 2018 Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: README.md ================================================ # Frozen Graph TensorFlow Lei Mao ## Introduction This repository has the examples of saving, loading, and running inference for frozen graph in TensorFlow 1.x and 2.x. ## Files ``` . ├── LICENSE.md ├── README.md ├── TensorFlow_v1 │   ├── cifar.py │   ├── cnn.py │   ├── inspect_signature.py │   ├── main.py │   ├── README.md │   ├── test_pb.py │   └── utils.py └── TensorFlow_v2 ├── example_1.py ├── example_2.py ├── README.md └── utils.py ``` ## Blogs * [Save, Load and Inference From TensorFlow Frozen Graph](https://leimao.github.io/blog/Save-Load-Inference-From-TF-Frozen-Graph/) * [Save, Load and Inference From TensorFlow 2.x Frozen Graph](https://leimao.github.io/blog/Save-Load-Inference-From-TF2-Frozen-Graph/) ## Examples * [TensorFlow 1.x](https://github.com/leimao/Frozen_Graph_TensorFlow/tree/master/TensorFlow_v1) * [TensorFlow 2.x](https://github.com/leimao/Frozen_Graph_TensorFlow/tree/master/TensorFlow_v2) ================================================ FILE: TensorFlow_v1/.gitignore ================================================ __pycache__ frozen_models models ================================================ FILE: TensorFlow_v1/README.md ================================================ # Frozen Graph TensorFlow 1.x Lei Mao ## Introduction This repository was modified from my previous [simple CNN model](https://github.com/leimao/Convolutional_Neural_Network_CIFAR10) to classify CIFAR10 dataset. It consist training, saving model to frozen graph ``pb`` file, load ``pb`` file and do inference in TensorFlow. The tutorial with detailed description is available on my [blog](https://leimao.github.io/blog/Save-Load-Inference-From-TF-Frozen-Graph/). To the best of my knowledge, there is few similar tutorials on the internet. I wish this sample code could help you to prepare your own ``pb`` file for deployment. ## Dependencies * Python 3.6 * Numpy 1.14 * TensorFlow 1.12 * Matplotlib 2.1.1 (for demo purpose) ## Files ```bash . ├── cifar.py ├── cnn.py ├── main.py ├── README.md └── utils.py ``` ## Features * User-friendly CNN API wrapped * Allows changing learning rate and dropout rate in real time * No need for significant changes to codes in order to work for other tasks ## Usage ### Train and Test Model in TensorFlow ```bash $ python main.py --help usage: main.py [-h] [-train] [-test] [--lr LR] [--lr_decay LR_DECAY] [--dropout DROPOUT] [--batch_size BATCH_SIZE] [--epochs EPOCHS] [--optimizer OPTIMIZER] [--seed SEED] [--model_dir MODEL_DIR] [--model_filename MODEL_FILENAME] [--log_dir LOG_DIR] Train CNN on CIFAR10 dataset. optional arguments: -h, --help show this help message and exit -train, --train train model -test, --test test model --lr LR initial learning rate --lr_decay LR_DECAY learning rate decay --dropout DROPOUT dropout rate --batch_size BATCH_SIZE mini batch size --epochs EPOCHS number of epochs --optimizer OPTIMIZER optimizer --seed SEED random seed --model_dir MODEL_DIR model directory --model_filename MODEL_FILENAME model filename --log_dir LOG_DIR log directory ``` ```bash $ python main.py --train --test --epoch 30 --lr_decay 0.9 --dropout 0.5 ``` ### Test Model from PB File ```bash $ python test_pb.py --help usage: test_pb.py [-h] [--model_pb_filepath MODEL_PB_FILEPATH] Load and test model from frozen graph pb file. optional arguments: -h, --help show this help message and exit --model_pb_filepath MODEL_PB_FILEPATH model pb-format frozen graph file filepath ``` ```bash $ python test_pb.py ``` ## Update Log ### 2019/9/16 Replaced using the side effect of `tf.InteractiveSession` to set default graph for loading `graphdef` to using Python resource management `with` to set default graph for loading `graphdef`. ## Reference * [Save, Load and Inference From TensorFlow Frozen Graph](https://leimao.github.io/blog/Save-Load-Inference-From-TF-Frozen-Graph/) ================================================ FILE: TensorFlow_v1/cifar.py ================================================ import tensorflow as tf import numpy as np def train_test_split(x, y, train_fraction=0.9): # Split the data into training data and test data assert len(x) == len(y) dataset_size = len(x) idx = np.arange(len(x)) np.random.shuffle(idx) idx_split = int(dataset_size * train_fraction) x_train = x[:idx_split] y_train = y[:idx_split] x_test = x[idx_split:] y_test = y[idx_split:] return (x_train, y_train), (x_test, y_test) class CIFAR10(object): def __init__(self, train_fraction=0.9): (self.x_train, self.y_train), (self.x_test, self.y_test) = tf.keras.datasets.cifar10.load_data() (self.x_train, self.y_train), (self.x_valid, self.y_valid) = train_test_split( x=self.x_train, y=self.y_train, train_fraction=train_fraction) assert np.array_equal(np.unique(self.y_train), np.unique(self.y_test)) == True self.num_classes = len(np.unique(self.y_train)) self.input_size = list(self.x_train.shape[1:]) # Convert integer label to binary vector self.y_train_onehot = tf.keras.utils.to_categorical( self.y_train, self.num_classes) self.y_valid_onehot = tf.keras.utils.to_categorical( self.y_valid, self.num_classes) self.y_test_onehot = tf.keras.utils.to_categorical( self.y_test, self.num_classes) # Image scaling self.x_train = self.x_train.astype('float32') self.x_valid = self.x_valid.astype('float32') self.x_test = self.x_test.astype('float32') self.x_train /= 255 self.x_valid /= 255 self.x_test /= 255 ================================================ FILE: TensorFlow_v1/cnn.py ================================================ import tensorflow as tf import os from tensorflow.python.tools import freeze_graph from tensorflow.python.framework import graph_util from tensorflow.python.saved_model import builder as saved_model_builder from tensorflow.python.saved_model import signature_def_utils from tensorflow.python.saved_model import signature_constants from tensorflow.python.saved_model import tag_constants from tensorflow.python.saved_model import utils as saved_model_utils class CNN(object): def __init__(self, input_size, num_classes, optimizer): self.num_classes = num_classes self.input_size = input_size self.optimizer = optimizer self.learning_rate = tf.placeholder(tf.float32, shape=[], name='learning_rate') self.dropout_rate = tf.placeholder(tf.float32, shape=[], name='dropout_rate') self.input = tf.placeholder(tf.float32, [None] + self.input_size, name='input') self.label = tf.placeholder(tf.float32, [None, self.num_classes], name='label') self.output = self.network_initializer() self.loss = self.loss_initializer() self.optimization = self.optimizer_initializer() self.saver = tf.train.Saver() self.sess = tf.Session() self.sess.run(tf.global_variables_initializer()) def network(self, input, dropout_rate): conv1 = tf.layers.conv2d(inputs=input, filters=64, kernel_size=[3, 3], padding='same', activation=tf.nn.relu, name='conv1') conv2 = tf.layers.conv2d(inputs=conv1, filters=64, kernel_size=[3, 3], padding='same', activation=tf.nn.relu, name='conv2') pool1 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=[2, 2], name='pool1') pool1_dropout = tf.layers.dropout(inputs=pool1, rate=dropout_rate, training=True, name='pool1_dropout') conv3 = tf.layers.conv2d(inputs=pool1_dropout, filters=128, kernel_size=[3, 3], padding='same', activation=tf.nn.relu, name='conv3') conv4 = tf.layers.conv2d(inputs=conv3, filters=128, kernel_size=[3, 3], padding='same', activation=tf.nn.relu, name='conv4') pool2 = tf.layers.max_pooling2d(inputs=conv4, pool_size=[2, 2], strides=[2, 2], name='pool2') pool2_dropout = tf.layers.dropout(inputs=pool2, rate=dropout_rate, training=True, name='pool2_dropout') conv5 = tf.layers.conv2d(inputs=pool2_dropout, filters=256, kernel_size=[3, 3], padding='same', activation=tf.nn.relu, name='conv5') pool3 = tf.layers.max_pooling2d(inputs=conv5, pool_size=[2, 2], strides=[2, 2], name='pool3') pool3_dropout = tf.layers.dropout(inputs=pool3, rate=dropout_rate, training=True, name='pool3_dropout') flat = tf.layers.flatten(inputs=pool3_dropout, name='flat') fc1 = tf.layers.dense(inputs=flat, units=256, activation=tf.nn.relu, name='fc1') fc1_dropout = tf.layers.dropout(inputs=fc1, rate=dropout_rate, training=True, name='fc1_dropout') fc2 = tf.layers.dense(inputs=fc1_dropout, units=self.num_classes, activation=None, name='fc2') # Give output node a output = tf.identity(fc2, name='output') return output def network_initializer(self): with tf.variable_scope('cnn') as scope: ouput = self.network(input=self.input, dropout_rate=self.dropout_rate) return ouput def loss_initializer(self): with tf.variable_scope('loss') as scope: cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2( labels=self.label, logits=self.output, name='cross_entropy') cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy_mean') return cross_entropy_mean def optimizer_initializer(self): if self.optimizer == 'Adam': optimizer = tf.train.AdamOptimizer( learning_rate=self.learning_rate).minimize(self.loss) else: optimizer = tf.train.GradientDescentOptimizer( learning_rate=self.learning_rate).minimize(self.loss) return optimizer def train(self, data, label, learning_rate, dropout_rate): _, train_loss = self.sess.run( [self.optimization, self.loss], feed_dict={ self.input: data, self.label: label, self.learning_rate: learning_rate, self.dropout_rate: dropout_rate }) return train_loss def validate(self, data, label): output, validate_loss = self.sess.run([self.output, self.loss], feed_dict={ self.input: data, self.label: label, self.dropout_rate: 0.0 }) return output, validate_loss def test(self, data): output = self.sess.run(self.output, feed_dict={ self.input: data, self.dropout_rate: 0.0 }) return output def save(self, directory, filename): if not os.path.exists(directory): os.makedirs(directory) filepath = os.path.join(directory, filename + '.ckpt') self.saver.save(self.sess, filepath) return filepath def save_signature(self, directory): signature = signature_def_utils.build_signature_def( inputs={ 'input': saved_model_utils.build_tensor_info(self.input), 'dropout_rate': saved_model_utils.build_tensor_info(self.dropout_rate) }, outputs={ 'output': saved_model_utils.build_tensor_info(self.output) }, method_name=signature_constants.PREDICT_METHOD_NAME) signature_map = { signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature } model_builder = saved_model_builder.SavedModelBuilder(directory) model_builder.add_meta_graph_and_variables( self.sess, tags=[tag_constants.SERVING], signature_def_map=signature_map, clear_devices=True) model_builder.save(as_text=False) def save_as_pb(self, directory, filename): if not os.path.exists(directory): os.makedirs(directory) # Save check point for graph frozen later ckpt_filepath = self.save(directory=directory, filename=filename) pbtxt_filename = filename + '.pbtxt' pbtxt_filepath = os.path.join(directory, pbtxt_filename) pb_filepath = os.path.join(directory, filename + '.pb') # This will only save the graph but the variables will not be saved. # You have to freeze your model first. tf.train.write_graph(graph_or_graph_def=self.sess.graph_def, logdir=directory, name=pbtxt_filename, as_text=True) # Freeze graph # Method 1 freeze_graph.freeze_graph(input_graph=pbtxt_filepath, input_saver='', input_binary=False, input_checkpoint=ckpt_filepath, output_node_names='cnn/output', restore_op_name='save/restore_all', filename_tensor_name='save/Const:0', output_graph=pb_filepath, clear_devices=True, initializer_nodes='') # Method 2 ''' graph = tf.get_default_graph() input_graph_def = graph.as_graph_def() output_node_names = ['cnn/output'] output_graph_def = graph_util.convert_variables_to_constants(self.sess, input_graph_def, output_node_names) with tf.gfile.GFile(pb_filepath, 'wb') as f: f.write(output_graph_def.SerializeToString()) ''' return pb_filepath def load(self, filepath): if os.path.splitext(filepath)[1] != '.ckpt': filepath += '.ckpt' self.saver.restore(self.sess, filepath) ================================================ FILE: TensorFlow_v1/inspect_signature.py ================================================ import tensorflow as tf from tensorflow.python.saved_model import tag_constants def retrieve_model_data_info(saved_model_path): with tf.Session() as sess: graph = tf.Graph() with graph.as_default(): metagraph = tf.saved_model.loader.load(sess, [tag_constants.SERVING], saved_model_path) inputs_mapping = dict( metagraph.signature_def['serving_default'].inputs) outputs_mapping = dict( metagraph.signature_def['serving_default'].outputs) print("Print output mapping: ", outputs_mapping) print("Print input mapping: ", inputs_mapping) retrieve_model_data_info('./model/signature') ================================================ FILE: TensorFlow_v1/main.py ================================================ import os import argparse import tensorflow as tf import numpy as np from cnn import CNN from cifar import CIFAR10 from utils import plot_curve, model_accuracy def train(learning_rate, learning_rate_decay, dropout_rate, mini_batch_size, epochs, optimizer, random_seed, model_directory, model_filename, log_directory): np.random.seed(random_seed) if not os.path.exists(log_directory): os.makedirs(log_directory) # Load CIFAR10 dataset cifar10 = CIFAR10() x_train = cifar10.x_train y_train = cifar10.y_train y_train_onehot = cifar10.y_train_onehot x_valid = cifar10.x_valid y_valid = cifar10.y_valid y_valid_onehot = cifar10.y_valid_onehot num_classes = cifar10.num_classes input_size = cifar10.input_size print('CIFAR10 Input Image Size: {}'.format(input_size)) model = CNN(input_size=input_size, num_classes=num_classes, optimizer=optimizer) train_accuracy_log = list() valid_accuracy_log = list() train_loss_log = list() for epoch in range(epochs): print('Epoch: %d' % epoch) learning_rate *= learning_rate_decay # Prepare mini batches on train set shuffled_idx = np.arange(len(x_train)) np.random.shuffle(shuffled_idx) mini_batch_idx = [ shuffled_idx[k:k + mini_batch_size] for k in range(0, len(x_train), mini_batch_size) ] # Validate on validation set valid_prediction_onehot = model.test(data=x_valid) valid_prediction = np.argmax(valid_prediction_onehot, axis=1).reshape( (-1, 1)) valid_accuracy = model_accuracy(label=y_valid, prediction=valid_prediction) print('Validation Accuracy: %f' % valid_accuracy) valid_accuracy_log.append(valid_accuracy) # Train on train set for i, idx in enumerate(mini_batch_idx): train_loss = model.train(data=x_train[idx], label=y_train_onehot[idx], learning_rate=learning_rate, dropout_rate=dropout_rate) if i % 200 == 0: train_prediction_onehot = model.test(data=x_train[idx]) train_prediction = np.argmax(train_prediction_onehot, axis=1).reshape((-1, 1)) train_accuracy = model_accuracy(label=y_train[idx], prediction=train_prediction) print('Training Loss: %f, Training Accuracy: %f' % (train_loss, train_accuracy)) if i == 0: train_accuracy_log.append(train_accuracy) train_loss_log.append(train_loss) model.save(directory=model_directory, filename=model_filename) print('Trained model saved successfully') model.save_as_pb(directory=model_directory, filename=model_filename) print('Trained model saved as pb successfully') # The directory should not exist before calling this method signature_dir = os.path.join(model_directory, 'signature') assert (not os.path.exists(signature_dir)) model.save_signature(directory=signature_dir) print('Trained model with signature saved successfully') plot_curve(train_losses = train_loss_log, train_accuracies = train_accuracy_log, valid_accuracies = valid_accuracy_log, \ filename = os.path.join(log_directory, 'training_curve.png')) def test(model_file): tf.reset_default_graph() # Load CIFAR10 dataset cifar10 = CIFAR10() x_test = cifar10.x_test y_test = cifar10.y_test y_test_onehot = cifar10.y_test_onehot num_classes = cifar10.num_classes input_size = cifar10.input_size model = CNN(input_size=input_size, num_classes=num_classes, optimizer='Adam') model.load(filepath=model_file) test_prediction_onehot = model.test(data=x_test) test_prediction = np.argmax(test_prediction_onehot, axis=1).reshape( (-1, 1)) test_accuracy = model_accuracy(label=y_test, prediction=test_prediction) print('Test Accuracy: %f' % test_accuracy) def main(): # Default settings learning_rate_default = 0.001 learning_rate_decay_default = 0.9 dropout_rate_default = 0.5 mini_batch_size_default = 64 epochs_default = 30 optimizer_default = 'Adam' random_seed_default = 0 model_directory_default = 'model' model_filename_default = 'cifar10_cnn' log_directory_default = 'log' # Argparser parser = argparse.ArgumentParser( description='Train CNN on CIFAR10 dataset.') parser.add_argument('-train', '--train', help='train model', action='store_true') parser.add_argument('-test', '--test', help='test model', action='store_true') parser.add_argument('--lr', type=float, help='initial learning rate', default=learning_rate_default) parser.add_argument('--lr_decay', type=float, help='learning rate decay', default=learning_rate_decay_default) parser.add_argument('--dropout', type=float, help='dropout rate', default=dropout_rate_default) parser.add_argument('--batch_size', type=int, help='mini batch size', default=mini_batch_size_default) parser.add_argument('--epochs', type=int, help='number of epochs', default=epochs_default) parser.add_argument('--optimizer', type=str, help='optimizer', default=optimizer_default) parser.add_argument('--seed', type=int, help='random seed', default=random_seed_default) parser.add_argument('--model_dir', type=str, help='model directory', default=model_directory_default) parser.add_argument('--model_filename', type=str, help='model filename', default=model_filename_default) parser.add_argument('--log_dir', type=str, help='log directory', default=log_directory_default) argv = parser.parse_args() # Post-process argparser learning_rate = argv.lr learning_rate_decay = argv.lr_decay dropout_rate = argv.dropout mini_batch_size = argv.batch_size epochs = argv.epochs optimizer = argv.optimizer random_seed = argv.seed model_directory = argv.model_dir model_filename = argv.model_filename log_directory = argv.log_dir if argv.train: print('Training CNN on CIFAR10 dataset...') train(learning_rate=learning_rate, learning_rate_decay=learning_rate_decay, dropout_rate=dropout_rate, mini_batch_size=mini_batch_size, epochs=epochs, optimizer=optimizer, random_seed=random_seed, model_directory=model_directory, model_filename=model_filename, log_directory=log_directory) if argv.test: print('Testing CNN on CIFAR10 dataset...') test(model_file=os.path.join(model_directory_default, model_filename_default)) if __name__ == '__main__': main() ================================================ FILE: TensorFlow_v1/test_pb.py ================================================ import tensorflow as tf import numpy as np import argparse from cifar import CIFAR10 from utils import model_accuracy from tensorflow.python.framework import tensor_util # If load from pb, you may have to use get_tensor_by_name heavily. class CNN(object): def __init__(self, model_filepath): # The file path of model self.model_filepath = model_filepath # Initialize the model self.load_graph(model_filepath=self.model_filepath) def load_graph(self, model_filepath): ''' Lode trained model. ''' print('Loading model...') self.graph = tf.Graph() with tf.gfile.GFile(model_filepath, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) print('Check out the input placeholders:') nodes = [ n.name + ' => ' + n.op for n in graph_def.node if n.op in ('Placeholder') ] for node in nodes: print(node) with self.graph.as_default(): # Define input tensor self.input = tf.placeholder(np.float32, shape=[None, 32, 32, 3], name='input') self.dropout_rate = tf.placeholder(tf.float32, shape=[], name='dropout_rate') tf.import_graph_def(graph_def, { 'input': self.input, 'dropout_rate': self.dropout_rate }) self.graph.finalize() print('Model loading complete!') # Get layer names layers = [op.name for op in self.graph.get_operations()] for layer in layers: print(layer) """ # Check out the weights of the nodes weight_nodes = [n for n in graph_def.node if n.op == 'Const'] for n in weight_nodes: print("Name of the node - %s" % n.name) # print("Value - " ) # print(tensor_util.MakeNdarray(n.attr['value'].tensor)) """ # In this version, tf.InteractiveSession and tf.Session could be used interchangeably. # self.sess = tf.InteractiveSession(graph = self.graph) self.sess = tf.Session(graph=self.graph) def test(self, data): # Know your output node name output_tensor = self.graph.get_tensor_by_name("import/cnn/output:0") output = self.sess.run(output_tensor, feed_dict={ self.input: data, self.dropout_rate: 0 }) return output def test_from_frozen_graph(model_filepath): tf.reset_default_graph() # Load CIFAR10 dataset cifar10 = CIFAR10() x_test = cifar10.x_test y_test = cifar10.y_test y_test_onehot = cifar10.y_test_onehot num_classes = cifar10.num_classes input_size = cifar10.input_size # Test 500 samples x_test = x_test[0:500] y_test = y_test[0:500] model = CNN(model_filepath=model_filepath) test_prediction_onehot = model.test(data=x_test) test_prediction = np.argmax(test_prediction_onehot, axis=1).reshape( (-1, 1)) test_accuracy = model_accuracy(label=y_test, prediction=test_prediction) print('Test Accuracy: %f' % test_accuracy) def main(): model_pb_filepath_default = './model/cifar10_cnn.pb' # Argparser parser = argparse.ArgumentParser( description='Load and test model from frozen graph pb file.') parser.add_argument('--model_pb_filepath', type=str, help='model pb-format frozen graph file filepath', default=model_pb_filepath_default) argv = parser.parse_args() model_pb_filepath = argv.model_pb_filepath test_from_frozen_graph(model_filepath=model_pb_filepath) if __name__ == '__main__': main() ================================================ FILE: TensorFlow_v1/utils.py ================================================ import numpy as np import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt def model_accuracy(label, prediction): # Evaluate the trained model return np.sum(label == prediction) / len(prediction) def plot_curve(train_losses, train_accuracies, valid_accuracies, savefig=True, showfig=False, filename='training_curve.png'): x = np.arange(len(train_losses)) y1 = train_accuracies y2 = valid_accuracies y3 = train_losses fig, ax1 = plt.subplots(figsize=(12, 8)) ax2 = ax1.twinx() ax1.plot(x, y1, color='b', marker='o', label='Training Accuracy') ax1.plot(x, y2, color='g', marker='o', label='Validation Accuracy') ax2.plot(x, y3, color='r', marker='o', label='Training Loss') ax1.set_xlabel('Epochs') ax1.set_ylabel('Accuracy') ax2.set_ylabel('Loss') ax1.legend() ax2.legend() if savefig: fig.savefig(filename, format='png', dpi=600, bbox_inches='tight') if showfig: plt.show() plt.close() return ================================================ FILE: TensorFlow_v2/.gitignore ================================================ __pycache__ frozen_models models ================================================ FILE: TensorFlow_v2/README.md ================================================ # Frozen Graph TensorFlow 2.x Lei Mao ## Introduction TensorFlow 1.x provided interface to freeze models via `tf.Session`. However, since TensorFlow 2.x removed `tf.Session`, freezing models in TensorFlow 2.x had been a problem to most of the users. In this repository, several simple concrete examples have been implemented to demonstrate how to freeze models and run inference using frozen models in TensorFlow 2.x. The frozen models are also fully compatible with inference using TensorFlow 1.x, TensorFlow 2.x, ONNX Runtime, and TensorRT. ## Usages ### Docker Container We use TensorFlow 2.3 Docker container from DockerHub. To download the Docker image, please run the following command in the terminal. ```bash $ docker pull tensorflow/tensorflow:2.3.0-gpu ``` To start the Docker container, please run the following command in the terminal. ```bash $ docker run --gpus all -it --rm -v $(pwd):/mnt tensorflow/tensorflow:2.3.0-gpu ``` ### Examples #### Example 1 We would train a simple fully connected neural network to classify the Fashion MNIST data. The model would be saved as `SavedModel` in the `models/simple_model` directory for completeness. In addition, the model would also be frozen and saved as `simple_frozen_graph.pb` in the `frozen_models` directory. To train, save, export, and run inference for the model, please run the following command in the terminal. ```bash $ python example_1.py ``` #### Example 2 We would train a simple recurrent neural network that has multiple inputs and outputs using random data. The model would be saved as `SavedModel` in the `models/complex_model` directory for completeness. In addition, the model would also be frozen and saved as `complex_frozen_graph.pb` in the `frozen_models` directory. To train, save, export, and run inference for the model, please run the following command in the terminal. ```bash $ python example_2.py ``` ### Convert Frozen Graph to ONNX If TensorFlow 1.x and `tf2onnx` have been installed, the frozen graph could be converted to ONNX model using the following command. ```bash $ python -m tf2onnx.convert --input ./frozen_models/frozen_graph.pb --output model.onnx --outputs Identity:0 --inputs x:0 ``` ### Convert Frozen Graph to UFF The frozen graph could also be converted to UFF model for TensorRT using the following command. ```bash $ convert-to-uff frozen_graph.pb -t -O Identity -o frozen_graph.uff ``` TensorRT 6.0 Docker image could be pulled from [NVIDIA NGC](https://ngc.nvidia.com/). ```bash $ docker pull nvcr.io/nvidia/tensorrt:19.12-py3 ``` ## References * [Migrate from TensorFlow 1.x to 2.x](https://www.tensorflow.org/guide/migrate) ================================================ FILE: TensorFlow_v2/example_1.py ================================================ import tensorflow as tf from tensorflow import keras from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 import numpy as np from utils import get_fashion_mnist_data, wrap_frozen_graph def main(): tf.random.set_seed(seed=0) # Get data (train_images, train_labels), (test_images, test_labels) = get_fashion_mnist_data() # Create Keras model model = keras.Sequential(layers=[ keras.layers.InputLayer(input_shape=(28, 28), name="input"), keras.layers.Flatten(input_shape=(28, 28), name="flatten"), keras.layers.Dense(128, activation="relu", name="dense"), keras.layers.Dense(10, activation="softmax", name="output") ], name="FCN") # Print model architecture model.summary() # Compile model with optimizer model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"]) # Train model model.fit(x={"input": train_images}, y={"output": train_labels}, epochs=1) # Test model test_loss, test_acc = model.evaluate(x={"input": test_images}, y={"output": test_labels}, verbose=2) print("-" * 50) print("Test accuracy: ") print(test_acc) # Get predictions for test images predictions = model.predict(test_images) # Print the prediction for the first image print("-" * 50) print("Example TensorFlow prediction reference:") print(predictions[0]) # Save model to SavedModel format tf.saved_model.save(model, "./models/simple_model") # Convert Keras model to ConcreteFunction full_model = tf.function(lambda x: model(x)) full_model = full_model.get_concrete_function( x=tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype)) # Get frozen ConcreteFunction frozen_func = convert_variables_to_constants_v2(full_model) frozen_func.graph.as_graph_def() layers = [op.name for op in frozen_func.graph.get_operations()] print("-" * 50) print("Frozen model layers: ") for layer in layers: print(layer) print("-" * 50) print("Frozen model inputs: ") print(frozen_func.inputs) print("Frozen model outputs: ") print(frozen_func.outputs) # Save frozen graph from frozen ConcreteFunction to hard drive tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir="./frozen_models", name="simple_frozen_graph.pb", as_text=False) # Load frozen graph using TensorFlow 1.x functions with tf.io.gfile.GFile("./frozen_models/simple_frozen_graph.pb", "rb") as f: graph_def = tf.compat.v1.GraphDef() loaded = graph_def.ParseFromString(f.read()) # Wrap frozen graph to ConcreteFunctions frozen_func = wrap_frozen_graph(graph_def=graph_def, inputs=["x:0"], outputs=["Identity:0"], print_graph=True) print("-" * 50) print("Frozen model inputs: ") print(frozen_func.inputs) print("Frozen model outputs: ") print(frozen_func.outputs) # Get predictions for test images frozen_graph_predictions = frozen_func(x=tf.constant(test_images))[0] # Print the prediction for the first image print("-" * 50) print("Example TensorFlow frozen graph prediction reference:") print(frozen_graph_predictions[0].numpy()) # The two predictions should be almost the same. assert np.allclose(a=frozen_graph_predictions[0].numpy(), b=predictions[0], rtol=1e-05, atol=1e-08, equal_nan=False) if __name__ == "__main__": main() ================================================ FILE: TensorFlow_v2/example_2.py ================================================ import tensorflow as tf from tensorflow import keras from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 import numpy as np from utils import wrap_frozen_graph def main(): # Mysterious code # https://leimao.github.io/blog/TensorFlow-cuDNN-Failure/ gpu_devices = tf.config.experimental.list_physical_devices('GPU') for device in gpu_devices: tf.config.experimental.set_memory_growth(device, True) # Dummy example copied from TensorFlow # https://www.tensorflow.org/guide/keras/functional#models_with_multiple_inputs_and_outputs num_tags = 12 # Number of unique issue tags num_words = 10000 # Size of vocabulary obtained when preprocessing text data num_departments = 4 # Number of departments for predictions title_input = keras.Input( shape=(None,), name="title" ) # Variable-length sequence of ints body_input = keras.Input(shape=(None,), name="body") # Variable-length sequence of ints tags_input = keras.Input( shape=(num_tags,), name="tags" ) # Binary vectors of size `num_tags` # Embed each word in the title into a 64-dimensional vector title_features = keras.layers.Embedding(num_words, 64)(title_input) # Embed each word in the text into a 64-dimensional vector body_features = keras.layers.Embedding(num_words, 64)(body_input) # Reduce sequence of embedded words in the title into a single 128-dimensional vector title_features = keras.layers.LSTM(128)(title_features) # Reduce sequence of embedded words in the body into a single 32-dimensional vector body_features = keras.layers.LSTM(32)(body_features) # Merge all available features into a single large vector via concatenation x = keras.layers.concatenate([title_features, body_features, tags_input]) # Stick a logistic regression for priority prediction on top of the features priority_pred = keras.layers.Dense(1, name="priority")(x) # Stick a department classifier on top of the features department_pred = keras.layers.Dense(num_departments, name="department")(x) # Instantiate an end-to-end model predicting both priority and department model = keras.Model( inputs=[title_input, body_input, tags_input], outputs=[priority_pred, department_pred], ) model.compile( optimizer=keras.optimizers.RMSprop(1e-3), loss=[ keras.losses.BinaryCrossentropy(from_logits=True), keras.losses.CategoricalCrossentropy(from_logits=True), ], loss_weights=[1.0, 0.2], ) # Dummy input data title_data = np.random.randint(num_words, size=(1280, 10)).astype("float32") body_data = np.random.randint(num_words, size=(1280, 100)).astype("float32") tags_data = np.random.randint(2, size=(1280, num_tags)).astype("float32") # Dummy target data priority_targets = np.random.random(size=(1280, 1)) dept_targets = np.random.randint(2, size=(1280, num_departments)) model.fit( {"title": title_data, "body": body_data, "tags": tags_data}, {"priority": priority_targets, "department": dept_targets}, epochs=2, batch_size=32, ) predictions = model.predict({"title": title_data[0:1], "body": body_data[0:1], "tags": tags_data[0:1]}) predictions_priority = predictions[0] predictions_department = predictions[1] print("-" * 50) print("Example TensorFlow prediction reference:") print(predictions_priority) print(predictions_department) # Save model to SavedModel format tf.saved_model.save(model, "./models/complex_model") full_model = tf.function(lambda x: model(x)) full_model = full_model.get_concrete_function(x=(tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype), tf.TensorSpec(model.inputs[1].shape, model.inputs[1].dtype), tf.TensorSpec(model.inputs[2].shape, model.inputs[2].dtype))) # Get frozen ConcreteFunction # https://github.com/tensorflow/tensorflow/issues/36391#issuecomment-596055100 frozen_func = convert_variables_to_constants_v2(full_model, lower_control_flow=False) frozen_func.graph.as_graph_def() layers = [op.name for op in frozen_func.graph.get_operations()] print("-" * 50) print("Frozen model layers: ") for layer in layers: print(layer) print("-" * 50) print("Frozen model inputs: ") print(frozen_func.inputs) print("Frozen model outputs: ") print(frozen_func.outputs) # Save frozen graph from frozen ConcreteFunction to hard drive tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir="./frozen_models", name="complex_frozen_graph.pb", as_text=False) # Load frozen graph using TensorFlow 1.x functions with tf.io.gfile.GFile("./frozen_models/complex_frozen_graph.pb", "rb") as f: graph_def = tf.compat.v1.GraphDef() loaded = graph_def.ParseFromString(f.read()) # Wrap frozen graph to ConcreteFunctions frozen_func = wrap_frozen_graph(graph_def=graph_def, inputs=["x:0", "x_1:0", "x_2:0"], outputs=["Identity:0", "Identity_1:0"], print_graph=True) # Note that we only have "one" input and "output" for the loaded frozen function print("-" * 50) print("Frozen model inputs: ") print(frozen_func.inputs) print("Frozen model outputs: ") print(frozen_func.outputs) # Get predictions frozen_graph_predictions = frozen_func(x=tf.constant(title_data[0:1]), x_1=tf.constant(body_data[0:1]), x_2=tf.constant(tags_data[0:1])) frozen_graph_predictions_priority = frozen_graph_predictions[0] frozen_graph_predictions_department = frozen_graph_predictions[1] print("-" * 50) print("Example TensorFlow frozen graph prediction reference:") print(frozen_graph_predictions_priority.numpy()) print(frozen_graph_predictions_department.numpy()) # The two predictions should be almost the same. assert np.allclose(a=frozen_graph_predictions_priority.numpy(), b=predictions_priority, rtol=1e-05, atol=1e-08, equal_nan=False) assert np.allclose(a=frozen_graph_predictions_department.numpy(), b=predictions_department, rtol=1e-05, atol=1e-08, equal_nan=False) if __name__ == "__main__": main() ================================================ FILE: TensorFlow_v2/utils.py ================================================ import tensorflow as tf from tensorflow import keras import numpy as np def get_fashion_mnist_data(): fashion_mnist = keras.datasets.fashion_mnist (train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data() class_names = [ "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot" ] train_images = train_images.astype(np.float32) / 255.0 test_images = test_images.astype(np.float32) / 255.0 return (train_images, train_labels), (test_images, test_labels) def wrap_frozen_graph(graph_def, inputs, outputs, print_graph=False): def _imports_graph_def(): tf.compat.v1.import_graph_def(graph_def, name="") wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, []) import_graph = wrapped_import.graph if print_graph == True: print("-" * 50) print("Frozen model layers: ") layers = [op.name for op in import_graph.get_operations()] for layer in layers: print(layer) print("-" * 50) return wrapped_import.prune( tf.nest.map_structure(import_graph.as_graph_element, inputs), tf.nest.map_structure(import_graph.as_graph_element, outputs))