Useful Tensorflow 2.1.0 snippets

Useful Tensorflow 2.1.0 snippets.md

Useful Tensorflow 2.1.0 snippets

This is a collection of random things which are useful in Tensorflow 2.1.0, but not necessarily 100% intuitive from the documentation.

View available devices (including GPUs)

To view available devices, use the tf.config.list_physical_devices() command, as described here. The command accepts an optional device_type argument, which for example can be used to only list GPU devices, and returns a list of available devices, EG:

>>> tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
>>> tf.config.list_physical_devices()
[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

Disable TensorFlow logging

The following commands can be used before importing TensorFlow to disable TensorFlow logging, as described in this Stack Overflow answer (NB this applies to TensorFlow 1.X as well as 2.X):

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 
import tensorflow as tf

To only disable WARNING and INFO messages (but not ERROR messages), use os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' (see link above)

For some warning messages, it may also be necessary to include the following command after importing TensorFlow:

tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

Extracting a subgraph from a tf.function

Say we have a simple tf.function, such as the one below:

import tensorflow as tf

# Define Python function
def py_func(a, b, c, d):
    """
    g   = d + f
        = d + (c * e)
        = d + (c * (a + b))
    """
    print("Tracing")
    e = tf.add(a, b, name="first_add")
    f = tf.multiply(c, e, name="first_multiply")
    g = tf.add(d, f, name="second_add")
    return g

# Define tf.function and typical input
tf_func = tf.function(py_func)
a, b, c, d = [tf.constant(i + 1) for i in range(4)]
print(tf_func(a, b, c, d)) # 4 + (3 * (1 + 2)) = 13

Say we want to extract a subgraph from this tf.function, for example we want to set the result of first_add directly, provide the inputs c and d as normal, and then get the output from second_add. This might be useful, for example if we receive a saved Tensorflow model which we want to convert to TensorFlow Lite format, but the model contains preprocessing operations at the trunk which are unsupported in TensorFlow lite, so these unsupported operations need to be separated from the supported operations. One way to achieve this is by using the prune method of a WrappedFunction (returned by tf.compat.v1.wrap_function), and one way to create a WrappedFunction is to convert the tf.function into a graph_def and trace the graph_def's import back into Tensorflow:

# Create graph def
graph_def = tf_func.get_concrete_function(a, b, c, d).graph.as_graph_def()

# Prune model using graph def
wrapped_func = tf.compat.v1.wrap_function(
    lambda: tf.compat.v1.import_graph_def(graph_def, name=""), [])
op_output = lambda name: wrapped_func.graph.as_graph_element(name).outputs[0]
pruned_model = wrapped_func.prune(
    feeds=list(map(op_output, ["first_add", "c", "d"])),
    fetches=op_output("second_add"))

# Convert to tf.function
pruned_tf_func = tf.function(pruned_model)
first_add_result = tf.constant(2)
print(pruned_tf_func(first_add_result, c, d)) # 4 + (3 * 2) = 10

Console output:

Tracing
tf.Tensor(13, shape=(), dtype=int32)
tf.Tensor(10, shape=(), dtype=int32)

This Github issue may provide an alternative method for extracting a subgraph using the input_map method of the tf.import_graph_def function.

Saving a model

A function-object returned by tf.function cannot be saved just using tf.saved_model.save because the funciton-object is not "trackable". In order to save a tf.function, it should be added as an attribute of a trackable object such as a tf.Module; using the tf_func defined above:

to_export = tf.Module()
to_export.call = tf_func
tf.saved_model.save(to_export, export_dir)

restored_module = tf.saved_model.load(export_dir)
restored_func = restored_module.call
print(restored_func(a, b, c, d)) # 13

An equivalent approach is to subclass the tf.Module class (note that the tf.function which is defined must be called once in order to trace-compile a graph):

class MyModule(tf.Module):
    @tf.function
    def __call__(self, a, b, c, d): return tf_func(a, b, c, d)

# Instantiate the trackable object, and call once to trace-compile a graph
module_func = MyModule()
module_func(a, b, c, d)
tf.saved_model.save(module_func, export_dir)

Viewing all the ops in a graph

Given a tf.function, a straghtforward way to see all of the ops in the corresponding graph is to use the get_operations method of a corresponding tf.Graph, which can be obtained as a property of a concrete function for the tf.function:

print(tf_func.get_concrete_function(a, b, c, d).graph.get_operations())

Console output:

[<tf.Operation 'a' type=Placeholder>, <tf.Operation 'b' type=Placeholder>, <tf.Operation 'c' type=Placeholder>, <tf.Operation 
'd' type=Placeholder>, <tf.Operation 'first_add' type=Add>, <tf.Operation 'first_multiply' type=Mul>, <tf.Operation 'second_add' type=Add>, <tf.Operation 'Identity' type=Identity>]

Visualising a graph in TensorBoard

Visualising the graph of a tf.function is straightforward using the guide to Examining the TensorFlow Graph, from which the following snippet is adapted:

# Create TensorBoard summary to view model graph
logdir = os.path.join(".", "TensorBoard", "simple_model")
writer = tf.summary.create_file_writer(logdir)
tf.summary.trace_on(graph=True, profiler=True)
# Trace the function-call
tf_func(a, b, c, d)
# Write the summary to disk
with writer.as_default():
    tf.summary.trace_export(name="tf_func", step=0, profiler_outdir=logdir)

Note that this might not work if tf_func is called with inputs which don't lead to it getting trace-compiled; see the documentation for tf.function for more information. To see the output from TensorBoard, navigate to the TensorBoard directory, in a command prompt enter tensorboard --logdir ., and in a web-browser, go to http://localhost:6006/.

Converting to TensorFlow Lite

Creating a TensorFlow Lite model from a concrete function is straightforward, using the guides to getting started and the converter Python API:

# Convert the model to tflite format
concrete_func = tf_func.get_concrete_function(a, b, c, d)
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
tflite_model = converter.convert()

# Save the tflite file
tfl_filename = "simple_model.tflite"
with open(tfl_filename, "wb") as tfl_file: tfl_file.write(tflite_model)

From a standard TensorFlow environment, the TensorFlow Lite Interpreter class can simply be accessed as tf.lite.Interpreter. On embedded platforms and/or where disk space is limited, it is possible to just install the TensorFlow Lite interpreter (without the full TensorFlow package), as described in the Python quickstart guide for TensorFlow Lite, in which case the Interpreter class can be accessed as tflite_runtime.interpreter.Interpreter. A conditional import statement (which will work in either environment) might look like this:

try:
    # Try importing the small tflite_runtime module
    print("Trying to import tensorflow lite runtime...")
    from tflite_runtime.interpreter import Interpreter
except ModuleNotFoundError:
    # Try importing the full tensorflow module
    try:
        print("TFLite runtime not found; trying to import full tensorflow...")
        import tensorflow as tf
        Interpreter = tf.lite.Interpreter
    except ModuleNotFoundError:
        # Couldn't import either module
        raise RuntimeError("Could not import Tensorflow or Tensorflow Lite")

Once the Interpreter class is available, the model can be run with the inputs a, b, c, d defined earlier as follows:

# Load TFLite model and allocate tensors
interpreter = Interpreter(model_path=tfl_filename)
interpreter.allocate_tensors()

# Get input and output tensor details, which are both a list of dict objects
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Set the input tensors
for i, t in enumerate([a, b, c, d]):
    interpreter.set_tensor(input_details[i]['index'], t)

# Get the results and print
interpreter.invoke()
tflite_results = interpreter.get_tensor(output_details[0]['index'])

print(tflite_results) # 13

Loading a protobuf model

If a model is saved in an old protobuf format from TensorFlow 1.X, it can be accessed in TensorFlow 2 by first importing it as a graph_def:

# Read the model's graph definition from the protopuf file
with tf.compat.v1.gfile.FastGFile(graph_filename, 'rb') as graph_file:
    graph_def = tf.compat.v1.GraphDef()
    graph_def.ParseFromString(graph_file.read())
    tf.import_graph_def(graph_def, name='')

If the desired input and output layer names are known (which can be achieved by visualising the graph in TensorBoard, for example), the graph_def can be converted to a tf.function via tf.compat.v1.wrap_function (as described in the guide for migrating code from TensorFlow 1 to 2):

# Convert the graph def to a wrapped function
wrapped_import = tf.compat.v1.wrap_function(
    lambda: tf.compat.v1.import_graph_def(graph_def, name=""), [])
import_graph = wrapped_import.graph
wrapped_func_model = wrapped_import.prune(
    tf.nest.map_structure(import_graph.as_graph_element, input_layer_name),
    tf.nest.map_structure(import_graph.as_graph_element, output_layer_name))

# Convert the wrapped function to a concrete function
tf_func_model = tf.function(wrapped_func_model)