Full Code of leimao/Frozen_Graph_TensorFlow for AI

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))