Optimize Trained Models for Inference

Graph Transform Tool

Great Blog Post by Pete Warden from Google

Types of Optimizations

Remove training-only operations (checkpoint saving, drop out)
Strip out unused nodes
Remove debug operations
Fold batch normalization ops into weights (super cool)
Round weights
Quantize weights

Original Model (CPU)

File Size

WARNING: If this directory doesn't exist, you need to train the model in an earlier notebook!!



In [ ]:

    
%%bash

ls -l /root/models/optimize_me/linear/cpu/

Graph



In [ ]:

    
%%bash

summarize_graph --in_graph=/root/models/optimize_me/linear/cpu/unoptimized_cpu.pb



In [ ]:

    
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import re
from google.protobuf import text_format
from tensorflow.core.framework import graph_pb2

def convert_graph_to_dot(input_graph, output_dot, is_input_graph_binary):
    graph = graph_pb2.GraphDef()
    with open(input_graph, "rb") as fh:
        if is_input_graph_binary:
            graph.ParseFromString(fh.read())
        else:
            text_format.Merge(fh.read(), graph)
    with open(output_dot, "wt") as fh:
        print("digraph graphname {", file=fh)
        for node in graph.node:
            output_name = node.name
            print("  \"" + output_name + "\" [label=\"" + node.op + "\"];", file=fh)
            for input_full_name in node.input:
                parts = input_full_name.split(":")
                input_name = re.sub(r"^\^", "", parts[0])
                print("  \"" + input_name + "\" -> \"" + output_name + "\";", file=fh)
        print("}", file=fh)
        print("Created dot file '%s' for graph '%s'." % (output_dot, input_graph))



In [ ]:

    
input_graph='/root/models/optimize_me/linear/cpu/unoptimized_cpu.pb'
output_dot='/root/notebooks/unoptimized_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)



In [ ]:

    
%%bash

dot -T png /root/notebooks/unoptimized_cpu.dot \
    -o /root/notebooks/unoptimized_cpu.png > /tmp/a.out



In [ ]:

    
from IPython.display import Image

Image('/root/notebooks/unoptimized_cpu.png', width=1024, height=768)

Strip Unused Nodes



In [ ]:

    
%%bash

transform_graph \
    --in_graph=/root/models/optimize_me/linear/cpu/unoptimized_cpu.pb \
    --out_graph=/root/models/optimize_me/linear/cpu/strip_unused_optimized_cpu.pb \
    --inputs='x_observed' \
    --outputs='add' \
    --transforms='
strip_unused_nodes'



In [ ]:

    
%%bash

ls -l /root/models/optimize_me/linear/cpu

Graph



In [ ]:

    
%%bash

summarize_graph --in_graph=/root/models/optimize_me/linear/cpu/strip_unused_optimized_cpu.pb



In [ ]:

    
input_graph='/root/models/optimize_me/linear/cpu/strip_unused_optimized_cpu.pb'
output_dot='/root/notebooks/strip_unused_optimized_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)



In [ ]:

    
%%bash

dot -T png /root/notebooks/strip_unused_optimized_cpu.dot \
    -o /root/notebooks/strip_unused_optimized_cpu.png > /tmp/a.out



In [ ]:

    
from IPython.display import Image

Image('/root/notebooks/strip_unused_optimized_cpu.png')

Remove Nodes

Remove pesky Identity and CheckNumerics



In [ ]:

    
%%bash

transform_graph \
    --in_graph=/root/models/optimize_me/linear/cpu/strip_unused_optimized_cpu.pb \
    --out_graph=/root/models/optimize_me/linear/cpu/remove_nodes_optimized_cpu.pb \
    --inputs='x_observed' \
    --outputs='add' \
    --transforms='
strip_unused_nodes
remove_nodes(op=Identity, op=CheckNumerics)'



In [ ]:

    
%%bash

ls -l /root/models/optimize_me/linear/cpu

Graph



In [ ]:

    
%%bash

summarize_graph --in_graph=/root/models/optimize_me/linear/cpu/remove_nodes_optimized_cpu.pb



In [ ]:

    
input_graph='/root/models/optimize_me/linear/cpu/remove_nodes_optimized_cpu.pb'
output_dot='/root/notebooks/remove_nodes_optimized_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)



In [ ]:

    
%%bash

dot -T png /root/notebooks/remove_nodes_optimized_cpu.dot \
    -o /root/notebooks/remove_nodes_optimized_cpu.png > /tmp/a.out



In [ ]:

    
from IPython.display import Image

Image('/root/notebooks/remove_nodes_optimized_cpu.png')

Fold Constants



In [ ]:

    
%%bash

transform_graph \
--in_graph=/root/models/optimize_me/linear/cpu/unoptimized_cpu.pb \
--out_graph=/root/models/optimize_me/linear/cpu/fold_constants_optimized_cpu.pb \
--inputs='x_observed' \
--outputs='add' \
--transforms='
strip_unused_nodes
remove_nodes(op=Identity, op=CheckNumerics)
fold_constants(ignore_errors=true)'

File Size



In [ ]:

    
%%bash

ls -l /root/models/optimize_me/linear/cpu

Graph



In [ ]:

    
%%bash

summarize_graph --in_graph=/root/models/optimize_me/linear/cpu/fold_constants_optimized_cpu.pb



In [ ]:

    
input_graph='/root/models/optimize_me/linear/cpu/fold_constants_optimized_cpu.pb'
output_dot='/root/notebooks/fold_constants_optimized_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)



In [ ]:

    
%%bash

dot -T png /root/notebooks/fold_constants_optimized_cpu.dot \
    -o /root/notebooks/fold_constants_optimized_cpu.png > /tmp/a.out



In [ ]:

    
from IPython.display import Image

Image('/root/notebooks/fold_constants_optimized_cpu.png')

Fold Batch Normalizations

Prereq: fold_constants



In [ ]:

    
%%bash

transform_graph \
--in_graph=/root/models/optimize_me/linear/cpu/unoptimized_cpu.pb \
--out_graph=/root/models/optimize_me/linear/cpu/fold_batch_norms_optimized_cpu.pb \
--inputs='x_observed' \
--outputs='add' \
--transforms='
strip_unused_nodes
remove_nodes(op=Identity, op=CheckNumerics)
fold_constants(ignore_errors=true)
fold_batch_norms
fold_old_batch_norms'

File Size



In [ ]:

    
%%bash

ls -l /root/models/optimize_me/linear/cpu

Graph



In [ ]:

    
%%bash

summarize_graph --in_graph=/root/models/optimize_me/linear/cpu/fold_batch_norms_optimized_cpu.pb



In [ ]:

    
input_graph='/root/models/optimize_me/linear/cpu/fold_batch_norms_optimized_cpu.pb'
output_dot='/root/notebooks/fold_batch_norms_optimized_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)



In [ ]:

    
%%bash

dot -T png /root/notebooks/fold_batch_norms_optimized_cpu.dot \
    -o /root/notebooks/fold_batch_norms_optimized_cpu.png > /tmp/a.out



In [ ]:

    
from IPython.display import Image

Image('/root/notebooks/fold_batch_norms_optimized_cpu.png')

Quantize Weights

Prereq: fold_batch_norms



In [ ]:

    
%%bash

transform_graph \
--in_graph=/root/models/optimize_me/linear/cpu/unoptimized_cpu.pb \
--out_graph=/root/models/optimize_me/linear/cpu/quantize_weights_optimized_cpu.pb \
--inputs='x_observed' \
--outputs='add' \
--transforms='
strip_unused_nodes
remove_nodes(op=Identity, op=CheckNumerics)
fold_constants(ignore_errors=true)
fold_batch_norms
fold_old_batch_norms
quantize_weights'

File Size



In [ ]:

    
%%bash

ls -l /root/models/optimize_me/linear/cpu/

Graph



In [ ]:

    
%%bash

summarize_graph --in_graph=/root/models/optimize_me/linear/cpu/quantize_weights_optimized_cpu.pb



In [ ]:

    
input_graph='/root/models/optimize_me/linear/cpu/quantize_weights_optimized_cpu.pb'
output_dot='/root/notebooks/quantize_weights_optimized_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)



In [ ]:

    
%%bash

dot -T png /root/notebooks/quantize_weights_optimized_cpu.dot \
    -o /root/notebooks/quantize_weights_optimized_cpu.png > /tmp/a.out



In [ ]:

    
from IPython.display import Image

Image('/root/notebooks/quantize_weights_optimized_cpu.png')

Combine All Optimizations



In [ ]:

    
%%bash

transform_graph \
--in_graph=/root/models/optimize_me/linear/cpu/unoptimized_cpu.pb \
--out_graph=/root/models/optimize_me/linear/cpu/fully_optimized_cpu.pb \
--inputs='x_observed' \
--outputs='add' \
--transforms='
strip_unused_nodes
remove_nodes(op=Identity, op=CheckNumerics)
fold_constants(ignore_errors=true)
fold_batch_norms
fold_old_batch_norms
quantize_weights'

File Size



In [ ]:

    
%%bash

ls -l /root/models/optimize_me/linear/cpu/

Graph



In [ ]:

    
%%bash

summarize_graph --in_graph=/root/models/optimize_me/linear/cpu/fully_optimized_cpu.pb



In [ ]:

    
input_graph='/root/models/optimize_me/linear/cpu/fully_optimized_cpu.pb'
output_dot='/root/notebooks/fully_optimized_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)



In [ ]:

    
%%bash

dot -T png /root/notebooks/fully_optimized_cpu.dot \
    -o /root/notebooks/fully_optimized_cpu.png > /tmp/a.out



In [ ]:

    
from IPython.display import Image

Image('/root/notebooks/fully_optimized_cpu.png')