Tensorflow colocate issue with import_graph_def and ExponentialMovingAverage - tensorflow

Question

I had an issue with colocate information in the GraphDef.
Here are the high-level steps I follow
Train an Estimator using tf.train.ExponentialMovingAverage and use the EMA predictions for the PREDICT mode
Export to SavedModel
Reload the GraphDef from the SavedModel and remove unnecessary nodes with extract_sub_graph
Freeze the resulting graph (make variables into constants using checkpoint data) with freeze_graph_with_def_protos
At step 4. I get an error ValueError: Node 'layer/kernel/ExponentialMovingAverage' expects to be colocated with unknown node 'layer/kernel'
Here is the code I use to train the model
# train.py
import logging
from pathlib import Path
import sys
import tensorflow as tf
def ema_getter(ema):
def _ema_getter(getter, name, *args, **kwargs):
var = getter(name, *args, **kwargs)
ema_var = ema.average(var)
return ema_var if ema_var else var
return _ema_getter
def model_fn(features, labels, mode, params):
# pylint: disable=unused-argument
"""Dummy model_fn"""
if isinstance(features, dict): # For serving
features = features['feature']
predictions = tf.layers.dense(features, 1, name="layer")
predictions = tf.identity(predictions, name="predictions")
ema = tf.train.ExponentialMovingAverage(1.0)
variables = tf.get_collection(
tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name)
ema_op = ema.apply(variables)
with tf.variable_scope(tf.get_variable_scope(), reuse=True, custom_getter=ema_getter(ema)):
predictions_ema = tf.layers.dense(features, 1, name="layer")
predictions_ema = tf.identity(predictions_ema, name="predictions_ema")
if mode == tf.estimator.ModeKeys.PREDICT:
preds = {
"predictions_ema": predictions_ema
}
return tf.estimator.EstimatorSpec(mode, predictions=preds)
else:
loss = tf.nn.l2_loss(predictions - labels)
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(
mode, loss=loss)
elif mode == tf.estimator.ModeKeys.TRAIN:
train_op = tf.train.AdamOptimizer(learning_rate=0.5).minimize(
loss, global_step=tf.train.get_global_step())
return tf.estimator.EstimatorSpec(
mode, loss=loss, train_op=tf.group([train_op, ema_op]))
else:
raise NotImplementedError()
def train_generator_fn():
for number in range(100):
yield [number, number], [2 * number]
def train_input_fn():
shapes, types = (2, 1), (tf.float32, tf.float32)
dataset = tf.data.Dataset.from_generator(
train_generator_fn, output_types=types, output_shapes=shapes)
dataset = dataset.batch(20).repeat(200)
return dataset
def serving_input_receiver_fn():
"""Serving input_fn that builds features from placeholders
Returns
-------
tf.estimator.export.ServingInputReceiver
"""
number = tf.placeholder(dtype=tf.float32, shape=[None, 1], name='number')
receiver_tensors = {'number': number}
features = tf.tile(number, multiples=[1, 2])
return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
if __name__ == '__main__':
# Logging
Path('model').mkdir(exist_ok=True)
tf.logging.set_verbosity(logging.INFO)
handlers = [
logging.FileHandler('model/train.log'),
logging.StreamHandler(sys.stdout)
]
logging.getLogger('tensorflow').handlers = handlers
# Train estimator
estimator = tf.estimator.Estimator(model_fn, 'model', params={})
estimator.train(train_input_fn)
# Export
estimator.export_saved_model('saved_model', serving_input_receiver_fn)
and the code I use to optimize the graph
# optimize.py
from pathlib import Path
import tensorflow as tf
from tensorflow.python.tools.freeze_graph import freeze_graph_with_def_protos
from tensorflow.python.framework.graph_util import extract_sub_graph
from tensorflow.core.framework.graph_pb2 import GraphDef
def optimize_and_export(export_dir: str, output: str):
with tf.Session() as sess:
g = tf.saved_model.loader.load(sess, ["serve"], export_dir)
inference_graph = extract_sub_graph(g.graph_def, ["predictions_ema"])
g = freeze_graph_with_def_protos(
inference_graph,
None,
None,
"predictions_ema",
None,
None,
None,
None,
None,
input_saved_model_dir=export_dir,
saved_model_tags=["serve"],
)
tf.io.write_graph(g, logdir=str(Path(output).parent), name=Path(output).name, as_text=False)
if __name__ == '__main__':
export_dir = str(sorted(Path('saved_model').glob('*'))[0])
print(f"Reloading from {export_dir}")
optimize_and_export(export_dir, 'saved_model/final')

Answer 1

My understanding is that tf.train.ExponentialMovingAverage adds colocation information between the nodes (the orignal and the EMA version).
NB: colocation seems to mean "these variables should be located on the same device"
This colocate information is present in the graph protobuf (the SavedModel export).
When extracting the subgraph, only the EMA versions of the variables are kept, but the colocation information is preserved, which causes issues when creating the Graph, which tries to find the original colocated variables (not present anymore).
I found a way around by modifying the protobuf manually and removing all colocation information with
def optimize_and_export(export_dir: str, output: str):
with tf.Session() as sess:
g = tf.saved_model.loader.load(sess, ["serve"], export_dir)
inference_graph = extract_sub_graph(g.graph_def, ["predictions_ema"])
# Remove colocate information from GraphDef
for node in inference_graph.node:
if "_class" in node.attr:
del node.attr["_class"]
tf.logging.warning(f"Removing _class attr of {node.name}")
g = freeze_graph_with_def_protos(
inference_graph,
None,
None,
"predictions_ema",
None,
None,
None,
None,
None,
input_saved_model_dir=export_dir,
saved_model_tags=["serve"],
)
tf.io.write_graph(g, logdir=str(Path(output).parent), name=Path(output).name, as_text=False)