transform_graph.py

from tensorflow.tools.graph_transforms import TransformGraph

def transform_graph(graph_def, input_node_names, output_node_names):
    transforms = [
        "strip_unused_nodes",
        "remove_nodes(op=Identity, op=CheckNumerics)",
        "fold_constants(ignore_errors=true)",
        # remove colocation attribute
        "remove_attribute(attribute_name=_class)",
        "fold_batch_norms",
        "fold_old_batch_norms",
        "quantize_weights",
        "obfuscate_names"]
    transformed_graph_def = TransformGraph(
        graph_def, input_node_names, output_node_names, transforms)
    return transformed_graph_def