def _save_checkpoint_from_mock_model(self,
checkpoint_path,
use_moving_averages,
quantize=False,
num_channels=3):
g = tf.Graph()
with g.as_default():
mock_model = FakeModel()
inputs = tf.placeholder(tf.float32,
shape=[1, 10, 10, num_channels])
mock_model.predict(inputs, true_image_shapes=None)
if use_moving_averages:
tf.train.ExponentialMovingAverage(0.0).apply()
tf.train.get_or_create_global_step()
if quantize:
graph_rewriter_config = graph_rewriter_pb2.GraphRewriter()
graph_rewriter_config.quantization.delay = 500000
graph_rewriter_fn = graph_rewriter_builder.build(
graph_rewriter_config, is_training=False)
graph_rewriter_fn()
saver = tf.train.Saver()
init = tf.global_variables_initializer()
with self.test_session() as sess:
sess.run(init)
saver.save(sess, checkpoint_path)
def test_rewrite_nn_resize_op_quantized(self):
g = tf.Graph()
with g.as_default():
x = array_ops.placeholder(dtypes.float32, shape=(8, 10, 10, 8))
x_conv = tf.contrib.slim.conv2d(x, 8, 1)
y = array_ops.placeholder(dtypes.float32, shape=(8, 20, 20, 8))
s = ops.nearest_neighbor_upsampling(x_conv, 2)
t = s + y
graph_rewriter_config = graph_rewriter_pb2.GraphRewriter()
graph_rewriter_config.quantization.delay = 500000
graph_rewriter_fn = graph_rewriter_builder.build(
graph_rewriter_config, is_training=False)
graph_rewriter_fn()
exporter.rewrite_nn_resize_op(is_quantized=True)
resize_op_found = False
for op in g.get_operations():
if op.type == 'ResizeNearestNeighbor':
resize_op_found = True
self.assertEqual(op.inputs[0].op.type, 'FakeQuantWithMinMaxVars')
self.assertEqual(op.outputs[0].consumers()[0], t.op)
break
self.assertTrue(resize_op_found)
def get_graph_rewriter_config_from_file(graph_rewriter_config_file):
"""Parses config for graph rewriter.
Args:
graph_rewriter_config_file: file path to the graph rewriter config.
Returns:
graph_rewriter_pb2.GraphRewriter proto
"""
graph_rewriter_config = graph_rewriter_pb2.GraphRewriter()
with tf.gfile.GFile(graph_rewriter_config_file, "r") as f:
text_format.Merge(f.read(), graph_rewriter_config)
return graph_rewriter_config
def testQuantizationBuilderSetsUpCorrectEvalArguments(self):
with mock.patch.object(tf.contrib.quantize,
'experimental_create_eval_graph') as mock_quant_fn:
with mock.patch.object(tf.contrib.layers,
'summarize_collection') as mock_summarize_col:
graph_rewriter_proto = graph_rewriter_pb2.GraphRewriter()
graph_rewriter_proto.quantization.delay = 10
graph_rewrite_fn = graph_rewriter_builder.build(
graph_rewriter_proto, is_training=False)
graph_rewrite_fn()
_, kwargs = mock_quant_fn.call_args
self.assertEqual(kwargs['input_graph'], tf.get_default_graph())
mock_summarize_col.assert_called_with('quant_vars')
def testQuantizationBuilderSetsUpCorrectTrainArguments(self):
with mock.patch.object(
tf.contrib.quantize, 'create_training_graph') as mock_quant_fn:
with mock.patch.object(tf.contrib.layers,
'summarize_collection') as mock_summarize_col:
graph_rewriter_proto = graph_rewriter_pb2.GraphRewriter()
graph_rewriter_proto.quantization.delay = 10
graph_rewriter_proto.quantization.weight_bits = 8
graph_rewriter_proto.quantization.activation_bits = 8
graph_rewrite_fn = graph_rewriter_builder.build(
graph_rewriter_proto, is_training=True)
graph_rewrite_fn()
_, kwargs = mock_quant_fn.call_args
self.assertEqual(kwargs['input_graph'], tf.get_default_graph())
self.assertEqual(kwargs['quant_delay'], 10)
mock_summarize_col.assert_called_with('quant_vars')
def testQuantizationBuilderAddsQuantOverride(self):
graph = ops.Graph()
with graph.as_default():
self._buildGraph()
quant_overrides_proto = quant_overrides_pb2.QuantOverrides()
quant_config = quant_overrides_proto.quant_configs.add()
quant_config.op_name = 'test_graph/add_ab'
quant_config.quant_op_name = 'act_quant'
quant_config.fixed_range = True
quant_config.min = 0
quant_config.max = 6
quant_config.delay = 100
graph_rewriter_proto = graph_rewriter_pb2.GraphRewriter()
graph_rewriter_proto.quantization.delay = 10
graph_rewriter_proto.quantization.weight_bits = 8
graph_rewriter_proto.quantization.activation_bits = 8
graph_rewrite_fn = graph_rewriter_builder.build(
graph_rewriter_proto,
quant_overrides_config=quant_overrides_proto,
is_training=True)
graph_rewrite_fn()
act_quant_found = False
quant_delay_found = False
for op in graph.get_operations():
if (quant_config.quant_op_name in op.name
and op.type == 'FakeQuantWithMinMaxArgs'):
act_quant_found = True
min_val = op.get_attr('min')
max_val = op.get_attr('max')
self.assertEqual(min_val, quant_config.min)
self.assertEqual(max_val, quant_config.max)
if ('activate_quant' in op.name
and quant_config.quant_op_name in op.name
and op.type == 'Const'):
tensor = op.get_attr('value')
if tensor.int64_val[0] == quant_config.delay:
quant_delay_found = True
self.assertTrue(act_quant_found)
self.assertTrue(quant_delay_found)
