def test_export_model_with_detection_only_nodes(self): tmp_dir = self.get_temp_dir() trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt') self._save_checkpoint_from_mock_model(trained_checkpoint_prefix, use_moving_averages=True) output_directory = os.path.join(tmp_dir, 'output') inference_graph_path = os.path.join(output_directory, 'frozen_inference_graph.pb') with mock.patch.object( model_builder, 'build', autospec=True) as mock_builder: mock_builder.return_value = FakeModel(add_detection_masks=False) pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() exporter.export_inference_graph( input_type='image_tensor', pipeline_config=pipeline_config, trained_checkpoint_prefix=trained_checkpoint_prefix, output_directory=output_directory) inference_graph = self._load_inference_graph(inference_graph_path) with self.test_session(graph=inference_graph): inference_graph.get_tensor_by_name('image_tensor:0') inference_graph.get_tensor_by_name('detection_boxes:0') inference_graph.get_tensor_by_name('detection_scores:0') inference_graph.get_tensor_by_name('detection_classes:0') inference_graph.get_tensor_by_name('num_detections:0') with self.assertRaises(KeyError): inference_graph.get_tensor_by_name('detection_keypoints:0') inference_graph.get_tensor_by_name('detection_masks:0')
def test_export_model_with_quantization_nodes(self): tmp_dir = self.get_temp_dir() trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt') self._save_checkpoint_from_mock_model( trained_checkpoint_prefix, use_moving_averages=False, enable_quantization=True) output_directory = os.path.join(tmp_dir, 'output') inference_graph_path = os.path.join(output_directory, 'inference_graph.pbtxt') with mock.patch.object( model_builder, 'build', autospec=True) as mock_builder: mock_builder.return_value = FakeModel() pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() text_format.Merge( """graph_rewriter { quantization { delay: 50000 activation_bits: 8 weight_bits: 8 } }""", pipeline_config) exporter.export_inference_graph( input_type='image_tensor', pipeline_config=pipeline_config, trained_checkpoint_prefix=trained_checkpoint_prefix, output_directory=output_directory, write_inference_graph=True) self._load_inference_graph(inference_graph_path, is_binary=False) has_quant_nodes = False for v in tf.global_variables(): if v.op.name.endswith('act_quant/min'): has_quant_nodes = True break self.assertTrue(has_quant_nodes)
def test_export_graph_with_fixed_size_image_tensor_input(self): input_shape = [1, 320, 320, 3] tmp_dir = self.get_temp_dir() trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt') self._save_checkpoint_from_mock_model( trained_checkpoint_prefix, use_moving_averages=False) with mock.patch.object( model_builder, 'build', autospec=True) as mock_builder: mock_builder.return_value = FakeModel() output_directory = os.path.join(tmp_dir, 'output') pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.eval_config.use_moving_averages = False exporter.export_inference_graph( input_type='image_tensor', pipeline_config=pipeline_config, trained_checkpoint_prefix=trained_checkpoint_prefix, output_directory=output_directory, input_shape=input_shape) saved_model_path = os.path.join(output_directory, 'saved_model') self.assertTrue( os.path.exists(os.path.join(saved_model_path, 'saved_model.pb'))) with tf.Graph().as_default() as od_graph: with self.test_session(graph=od_graph) as sess: meta_graph = tf.saved_model.loader.load( sess, [tf.saved_model.tag_constants.SERVING], saved_model_path) signature = meta_graph.signature_def['serving_default'] input_tensor_name = signature.inputs['inputs'].name image_tensor = od_graph.get_tensor_by_name(input_tensor_name) self.assertSequenceEqual(image_tensor.get_shape().as_list(), input_shape)
def test_export_saved_model_and_run_inference(self): tmp_dir = self.get_temp_dir() trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt') self._save_checkpoint_from_mock_model(trained_checkpoint_prefix, use_moving_averages=False) output_directory = os.path.join(tmp_dir, 'output') saved_model_path = os.path.join(output_directory, 'saved_model') with mock.patch.object( model_builder, 'build', autospec=True) as mock_builder: mock_builder.return_value = FakeModel( add_detection_keypoints=True, add_detection_masks=True) pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.eval_config.use_moving_averages = False exporter.export_inference_graph( input_type='tf_example', pipeline_config=pipeline_config, trained_checkpoint_prefix=trained_checkpoint_prefix, output_directory=output_directory) tf_example_np = np.hstack([self._create_tf_example( np.ones((4, 4, 3)).astype(np.uint8))] * 2) with tf.Graph().as_default() as od_graph: with self.test_session(graph=od_graph) as sess: meta_graph = tf.saved_model.loader.load( sess, [tf.saved_model.tag_constants.SERVING], saved_model_path) signature = meta_graph.signature_def['serving_default'] input_tensor_name = signature.inputs['inputs'].name tf_example = od_graph.get_tensor_by_name(input_tensor_name) boxes = od_graph.get_tensor_by_name( signature.outputs['detection_boxes'].name) scores = od_graph.get_tensor_by_name( signature.outputs['detection_scores'].name) classes = od_graph.get_tensor_by_name( signature.outputs['detection_classes'].name) keypoints = od_graph.get_tensor_by_name( signature.outputs['detection_keypoints'].name) masks = od_graph.get_tensor_by_name( signature.outputs['detection_masks'].name) num_detections = od_graph.get_tensor_by_name( signature.outputs['num_detections'].name) (boxes_np, scores_np, classes_np, keypoints_np, masks_np, num_detections_np) = sess.run( [boxes, scores, classes, keypoints, masks, num_detections], feed_dict={tf_example: tf_example_np}) self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 0.8, 0.8]], [[0.5, 0.5, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0]]]) self.assertAllClose(scores_np, [[0.7, 0.6], [0.9, 0.0]]) self.assertAllClose(classes_np, [[1, 2], [2, 1]]) self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2])) self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4])) self.assertAllClose(num_detections_np, [2, 1])
def test_export_and_run_inference_with_encoded_image_string_tensor(self): tmp_dir = self.get_temp_dir() trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt') self._save_checkpoint_from_mock_model(trained_checkpoint_prefix, use_moving_averages=True) output_directory = os.path.join(tmp_dir, 'output') inference_graph_path = os.path.join(output_directory, 'frozen_inference_graph.pb') with mock.patch.object( model_builder, 'build', autospec=True) as mock_builder: mock_builder.return_value = FakeModel( add_detection_keypoints=True, add_detection_masks=True) pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.eval_config.use_moving_averages = False exporter.export_inference_graph( input_type='encoded_image_string_tensor', pipeline_config=pipeline_config, trained_checkpoint_prefix=trained_checkpoint_prefix, output_directory=output_directory) inference_graph = self._load_inference_graph(inference_graph_path) jpg_image_str = self._create_encoded_image_string( np.ones((4, 4, 3)).astype(np.uint8), 'jpg') png_image_str = self._create_encoded_image_string( np.ones((4, 4, 3)).astype(np.uint8), 'png') with self.test_session(graph=inference_graph) as sess: image_str_tensor = inference_graph.get_tensor_by_name( 'encoded_image_string_tensor:0') boxes = inference_graph.get_tensor_by_name('detection_boxes:0') scores = inference_graph.get_tensor_by_name('detection_scores:0') classes = inference_graph.get_tensor_by_name('detection_classes:0') keypoints = inference_graph.get_tensor_by_name('detection_keypoints:0') masks = inference_graph.get_tensor_by_name('detection_masks:0') num_detections = inference_graph.get_tensor_by_name('num_detections:0') for image_str in [jpg_image_str, png_image_str]: image_str_batch_np = np.hstack([image_str]* 2) (boxes_np, scores_np, classes_np, keypoints_np, masks_np, num_detections_np) = sess.run( [boxes, scores, classes, keypoints, masks, num_detections], feed_dict={image_str_tensor: image_str_batch_np}) self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 0.8, 0.8]], [[0.5, 0.5, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0]]]) self.assertAllClose(scores_np, [[0.7, 0.6], [0.9, 0.0]]) self.assertAllClose(classes_np, [[1, 2], [2, 1]]) self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2])) self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4])) self.assertAllClose(num_detections_np, [2, 1])
def main(_): pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() with tf.gfile.GFile(FLAGS.pipeline_config_path, 'r') as f: text_format.Merge(f.read(), pipeline_config) text_format.Merge(FLAGS.config_override, pipeline_config) if FLAGS.input_shape: input_shape = [ int(dim) if dim != '-1' else None for dim in FLAGS.input_shape.split(',') ] else: input_shape = None exporter.export_inference_graph( FLAGS.input_type, pipeline_config, FLAGS.trained_checkpoint_prefix, FLAGS.output_directory, input_shape=input_shape, write_inference_graph=FLAGS.write_inference_graph)
def test_export_graph_with_encoded_image_string_input(self): tmp_dir = self.get_temp_dir() trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt') self._save_checkpoint_from_mock_model(trained_checkpoint_prefix, use_moving_averages=False) with mock.patch.object( model_builder, 'build', autospec=True) as mock_builder: mock_builder.return_value = FakeModel() output_directory = os.path.join(tmp_dir, 'output') pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.eval_config.use_moving_averages = False exporter.export_inference_graph( input_type='encoded_image_string_tensor', pipeline_config=pipeline_config, trained_checkpoint_prefix=trained_checkpoint_prefix, output_directory=output_directory) self.assertTrue(os.path.exists(os.path.join( output_directory, 'saved_model', 'saved_model.pb')))
def test_raise_runtime_error_on_images_with_different_sizes(self): tmp_dir = self.get_temp_dir() trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt') self._save_checkpoint_from_mock_model(trained_checkpoint_prefix, use_moving_averages=True) output_directory = os.path.join(tmp_dir, 'output') inference_graph_path = os.path.join(output_directory, 'frozen_inference_graph.pb') with mock.patch.object( model_builder, 'build', autospec=True) as mock_builder: mock_builder.return_value = FakeModel( add_detection_keypoints=True, add_detection_masks=True) pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.eval_config.use_moving_averages = False exporter.export_inference_graph( input_type='encoded_image_string_tensor', pipeline_config=pipeline_config, trained_checkpoint_prefix=trained_checkpoint_prefix, output_directory=output_directory) inference_graph = self._load_inference_graph(inference_graph_path) large_image = self._create_encoded_image_string( np.ones((4, 4, 3)).astype(np.uint8), 'jpg') small_image = self._create_encoded_image_string( np.ones((2, 2, 3)).astype(np.uint8), 'jpg') image_str_batch_np = np.hstack([large_image, small_image]) with self.test_session(graph=inference_graph) as sess: image_str_tensor = inference_graph.get_tensor_by_name( 'encoded_image_string_tensor:0') boxes = inference_graph.get_tensor_by_name('detection_boxes:0') scores = inference_graph.get_tensor_by_name('detection_scores:0') classes = inference_graph.get_tensor_by_name('detection_classes:0') keypoints = inference_graph.get_tensor_by_name('detection_keypoints:0') masks = inference_graph.get_tensor_by_name('detection_masks:0') num_detections = inference_graph.get_tensor_by_name('num_detections:0') with self.assertRaisesRegexp(tf.errors.InvalidArgumentError, 'TensorArray.*shape'): sess.run( [boxes, scores, classes, keypoints, masks, num_detections], feed_dict={image_str_tensor: image_str_batch_np})
def test_export_graph_with_moving_averages(self): tmp_dir = self.get_temp_dir() trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt') self._save_checkpoint_from_mock_model(trained_checkpoint_prefix, use_moving_averages=True) output_directory = os.path.join(tmp_dir, 'output') with mock.patch.object( model_builder, 'build', autospec=True) as mock_builder: mock_builder.return_value = FakeModel() pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() pipeline_config.eval_config.use_moving_averages = True exporter.export_inference_graph( input_type='image_tensor', pipeline_config=pipeline_config, trained_checkpoint_prefix=trained_checkpoint_prefix, output_directory=output_directory) self.assertTrue(os.path.exists(os.path.join( output_directory, 'saved_model', 'saved_model.pb'))) expected_variables = set(['conv2d/bias', 'conv2d/kernel', 'global_step']) actual_variables = set( [var_name for var_name, _ in tf.train.list_variables(output_directory)]) self.assertTrue(expected_variables.issubset(actual_variables))
def test_export_graph_saves_pipeline_file(self): tmp_dir = self.get_temp_dir() trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt') self._save_checkpoint_from_mock_model(trained_checkpoint_prefix, use_moving_averages=True) output_directory = os.path.join(tmp_dir, 'output') with mock.patch.object( model_builder, 'build', autospec=True) as mock_builder: mock_builder.return_value = FakeModel() pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() exporter.export_inference_graph( input_type='image_tensor', pipeline_config=pipeline_config, trained_checkpoint_prefix=trained_checkpoint_prefix, output_directory=output_directory) expected_pipeline_path = os.path.join( output_directory, 'pipeline.config') self.assertTrue(os.path.exists(expected_pipeline_path)) written_pipeline_config = pipeline_pb2.TrainEvalPipelineConfig() with tf.gfile.GFile(expected_pipeline_path, 'r') as f: proto_str = f.read() text_format.Merge(proto_str, written_pipeline_config) self.assertProtoEquals(pipeline_config, written_pipeline_config)