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)