Using AIMET Tensorflow APIs with Keras Models

Introduction

Currently AIMET APIs support Tensorflow sessions. This example code shows a method for how to use AIMET if you have a Keras model by invoking AIMET on the back-end session and converting the returned session to a Keras model.

APIs

The method involves performing four steps. The steps are:

Step 1: Save the session returned by AIMET

aimet_tensorflow.utils.convert_tf_sess_to_keras.save_tf_session_single_gpu(sess, path, input_tensor, output_tensor)[source]

Saves TF session, meta graph and variables in the provided path

Parameters

  • sess (Session) – Input: tf.compat.v1.Session

  • path (str) – Path to save the session

  • input_tensor (str) – Name of starting op to the given graph

  • output_tensor (str) – Name of output op of the graph

Returns

None


Step 2: Model subclassing to load the corresponding session to Keras model

aimet_tensorflow.utils.convert_tf_sess_to_keras.load_tf_sess_variables_to_keras_single_gpu(path, compressed_ops)[source]

Creates a Keras model subclass and loads the saved session, meta graph and variables to Keras model

Parameters

  • path (str) – Path to load the tf session saved using save_session_graph_and_variables

  • compressed_ops (List[str]) – List of ops names skipped in Keras model creations. These are the the ops that AIMET compressed and are isolated from rest of the graph.

Return type

Model

Returns

Subclassed Keras Model


After these two steps, model can be used for single gpu training. For multi-gpu training, the next two steps needs to be followed:

Step 3: Saving the Keras model from step 2 to make it compatible with distribution strategy

aimet_tensorflow.utils.convert_tf_sess_to_keras.save_as_tf_module_multi_gpu(loading_path, saving_path, compressed_ops, input_shape)[source]

Loads a Keras model and re-saves the loaded object in the form of tf.Module

Parameters

  • loading_path (str) – Path to load the Keras Model

  • saving_path (str) – Path to save the object

  • compressed_ops (List[str]) – List of ops names for which we need to skip in Keras model creation. These are the the ops that AIMET compressed and are isolated from rest of the graph.

  • input_shape (Tuple) – shape of input to the model

Returns

None


Step 4: Model subclassing to load the corresponding Keras model

aimet_tensorflow.utils.convert_tf_sess_to_keras.load_keras_model_multi_gpu(loading_path, input_shape)[source]

This function loads the Keras model back, which can be used for funetuning within a strategy

Parameters

  • loading_path (str) – Path to load the Keras Model

  • input_shape (List) – the shape of stating tensor in graph ; for instance (224,224,3) for ResNet50 and MoblinetV1

Returns

subclassed Keras model


Code Example

Required imports

import tensorflow as tf
from aimet_tensorflow.utils.convert_tf_sess_to_keras import save_tf_session_single_gpu, save_as_tf_module_multi_gpu, \
    load_tf_sess_variables_to_keras_single_gpu, load_keras_model_multi_gpu

Steps to convert a TF session found after compression to Keras model

def convert_tf_session_to_keras_model():
    """
    Convert an AIMET  spatial SVD compressed session to a Keras model and train the Keras model with MirroredStrategy
    """
    sess = get_sess_from_keras_model()

    # For instance, if the first conv layer in MobilNetV1 graph is compressed, then:
    compressed_ops = ['conv1/Conv2D']
    compressed_sess = compress_session(sess, compressed_ops)

    # Defining the input and output convs of the session for MobileNet model
    input_op_name, output_op_name = "input_1:0", "act_softmax/Softmax:0"

    # Step 1: Single Saving the compressed session
    path = './saved_model_single_gpu'
    save_tf_session_single_gpu(compressed_sess, path, input_op_name, output_op_name)
    tf.keras.backend.clear_session()

    # Step 2: Loading the correspnding Keras Model
    tf.keras.backend.set_learning_phase(1)
    model = load_tf_sess_variables_to_keras_single_gpu(path, compressed_ops)

    # Single GPU training of the loaded Keras Model
    train(model)

    # To be able to do multi-gpu training the next two steps needs to be followed:
    # Step 3: Re-Saving the Keras model to make it compatible with distribution strategy
    saving_path = './saved_model_multi_gpu'
    save_as_tf_module_multi_gpu(path, saving_path, compressed_ops, input_shape=(224, 224, 3))

    tf.keras.backend.clear_session()

    with tf.distribute.MirroredStrategy().scope():
        tf.keras.backend.set_learning_phase(1)
        # Step 4: Loading the keras model and  Multi gpu training the model on given dataset
        model = load_keras_model_multi_gpu(saving_path, input_shape=[224, 224, 3])
        # Train model on Multi-GPU
        train(model)

Utility Functions

Required imports

import tensorflow as tf
from tensorflow.keras.applications import MobileNet
from keras.applications.vgg16 import preprocess_input

import numpy as np

from aimet_common.defs import CompressionScheme, CostMetric
from aimet_tensorflow.defs import SpatialSvdParameters
from aimet_tensorflow.compress import ModelCompressor
from aimet_tensorflow.defs import ModuleCompRatioPair

Utility function to get session from Keras model

def get_sess_from_keras_model():
    """
    Gets TF session from keras model
    :return: TF session
    """
    tf.keras.backend.clear_session()
    tf.keras.backend.set_learning_phase(1)
    _ = MobileNet(weights=None, input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()
    return sess

Utility function to get a compressed session

def compress_session(sess, compressible_ops):
    """
    Compressed TF session
    :param sess: Tf session
    :param compressible_ops: layers to compress
    :return: compressed session
    """
    layer_a = sess.graph.get_operation_by_name(compressible_ops[0])
    list_of_module_comp_ratio_pairs = [ModuleCompRatioPair(layer_a, 0.5)]
    manual_params = SpatialSvdParameters.ManualModeParams(
        list_of_module_comp_ratio_pairs=list_of_module_comp_ratio_pairs)
    params = SpatialSvdParameters(input_op_names=['input_1'], output_op_names=['act_softmax/Softmax'],
                                  mode=SpatialSvdParameters.Mode.manual, params=manual_params)
    scheme = CompressionScheme.spatial_svd
    metric = CostMetric.mac

    # pylint: disable=unused-argument
    def evaluate(sess, iterations, use_cuda):
        return 1

    sess, _ = ModelCompressor.compress_model(sess=sess,
                                             working_dir="./",
                                             eval_callback=evaluate,
                                             eval_iterations=None,
                                             input_shape=(1, 3, 224, 224),
                                             compress_scheme=scheme,
                                             cost_metric=metric,
                                             parameters=params)
    return sess

Utility function for training

def train(model):
    """
    Trains using fake dataset
    :param model: Keras model
    :return: trained model
    """
    # Create a fake dataset
    x_train = np.random.rand(32, 224, 224, 3)
    y_train = np.random.rand(32, )
    x_train = preprocess_input(x_train)
    y_train = tf.keras.utils.to_categorical(y_train, 1000)

    model.compile('rmsprop', 'mse')
    model.fit(x_train, y_train, epochs=1, batch_size=1, shuffle=False)
    return model