Skip to content

Commit

Permalink
No public description
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 639040818
  • Loading branch information
@tensorflower-gardener
tensorflower-gardener committed May 31, 2024
1 parent 477fa84 commit c9d1883
Show file tree
Hide file tree
Showing 2 changed files with 109 additions and 61 deletions.
95 changes: 63 additions & 32 deletions official/vision/serving/detection.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,27 +41,37 @@ def _padded_size(self):
return self._input_image_size

def _build_model(self):

nms_versions_supporting_dynamic_batch_size = {'batched', 'v2', 'v3'}
nms_version = self.params.task.model.detection_generator.nms_version
if (self._batch_size is None and
nms_version not in nms_versions_supporting_dynamic_batch_size):
logging.info('nms_version is set to `batched` because `%s` '
'does not support with dynamic batch size.', nms_version)
if (
self._batch_size is None
and nms_version not in nms_versions_supporting_dynamic_batch_size
):
logging.info(
'nms_version is set to `batched` because `%s` '
'does not support with dynamic batch size.',
nms_version,
)
self.params.task.model.detection_generator.nms_version = 'batched'

input_specs = tf_keras.layers.InputSpec(shape=[
self._batch_size, *self._padded_size, 3])
input_specs = tf_keras.layers.InputSpec(
shape=[self._batch_size, *self._padded_size, 3]
)

if isinstance(self.params.task.model, configs.maskrcnn.MaskRCNN):
model = factory.build_maskrcnn(
input_specs=input_specs, model_config=self.params.task.model)
input_specs=input_specs, model_config=self.params.task.model
)
elif isinstance(self.params.task.model, configs.retinanet.RetinaNet):
model = factory.build_retinanet(
input_specs=input_specs, model_config=self.params.task.model)
input_specs=input_specs, model_config=self.params.task.model
)
else:
raise ValueError('Detection module not implemented for {} model.'.format(
type(self.params.task.model)))
raise ValueError(
'Detection module not implemented for {} model.'.format(
type(self.params.task.model)
)
)

return model

Expand All @@ -73,7 +83,8 @@ def _build_anchor_boxes(self):
max_level=model_params.max_level,
num_scales=model_params.anchor.num_scales,
aspect_ratios=model_params.anchor.aspect_ratios,
anchor_size=model_params.anchor.anchor_size)
anchor_size=model_params.anchor.anchor_size,
)
return input_anchor(image_size=self._padded_size)

def _build_inputs(self, image):
Expand All @@ -85,7 +96,8 @@ def _build_inputs(self, image):

# Normalizes image with mean and std pixel values.
image = preprocess_ops.normalize_image(
image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)
image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB
)

image, image_info = preprocess_ops.resize_and_crop_image(
image,
Expand Down Expand Up @@ -131,20 +143,24 @@ def preprocess(
Args:
images: The images tensor.
Returns:
images: The images tensor cast to float.
anchor_boxes: Dict mapping anchor levels to anchor boxes.
image_info: Tensor containing the details of the image resizing.
"""
model_params = self.params.task.model
with tf.device('cpu:0'):
# Tensor Specs for map_fn outputs (images, anchor_boxes, and image_info).
images_spec = tf.TensorSpec(shape=self._padded_size + [3],
dtype=tf.float32)
images_spec = tf.TensorSpec(
shape=self._padded_size + [3], dtype=tf.float32
)

num_anchors = model_params.anchor.num_scales * len(
model_params.anchor.aspect_ratios) * 4
num_anchors = (
model_params.anchor.num_scales
* len(model_params.anchor.aspect_ratios)
* 4
)
anchor_shapes = []
for level in range(model_params.min_level, model_params.max_level + 1):
anchor_level_spec = tf.TensorSpec(
Expand All @@ -153,7 +169,8 @@ def preprocess(
math.ceil(self._padded_size[1] / 2**level),
num_anchors,
],
dtype=tf.float32)
dtype=tf.float32,
)
anchor_shapes.append((str(level), anchor_level_spec))

image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)
Expand All @@ -163,9 +180,14 @@ def preprocess(
tf.map_fn(
self._build_inputs,
elems=images,
fn_output_signature=(images_spec, dict(anchor_shapes),
image_info_spec),
parallel_iterations=32))
fn_output_signature=(
images_spec,
dict(anchor_shapes),
image_info_spec,
),
parallel_iterations=32,
),
)

return images, anchor_boxes, image_info

Expand All @@ -174,6 +196,7 @@ def serve(self, images: tf.Tensor):
Args:
images: uint8 Tensor of shape [batch_size, None, None, 3]
Returns:
Tensor holding detection output logits.
"""
Expand All @@ -190,10 +213,15 @@ def serve(self, images: tf.Tensor):
# [desired_height, desired_width], [y_scale, x_scale],
# [y_offset, x_offset]]. When input_type is tflite, input image is
# supposed to be preprocessed already.
image_info = tf.convert_to_tensor([[
self._input_image_size, self._input_image_size, [1.0, 1.0], [0, 0]
]],
dtype=tf.float32)
image_info = tf.convert_to_tensor(
[[
self._input_image_size,
self._input_image_size,
[1.0, 1.0],
[0, 0],
]],
dtype=tf.float32,
)
input_image_shape = image_info[:, 1, :]

# To overcome keras.Model extra limitation to save a model with layers that
Expand Down Expand Up @@ -226,20 +254,23 @@ def serve(self, images: tf.Tensor):
# point outputs.
if export_config.cast_num_detections_to_float:
detections['num_detections'] = tf.cast(
detections['num_detections'], dtype=tf.float32)
detections['num_detections'], dtype=tf.float32
)
if export_config.cast_detection_classes_to_float:
detections['detection_classes'] = tf.cast(
detections['detection_classes'], dtype=tf.float32)
detections['detection_classes'], dtype=tf.float32
)

final_outputs = {
'detection_boxes': detections['detection_boxes'],
'detection_scores': detections['detection_scores'],
'detection_classes': detections['detection_classes'],
'num_detections': detections['num_detections']
'num_detections': detections['num_detections'],
}
if 'detection_outer_boxes' in detections:
final_outputs['detection_outer_boxes'] = (
detections['detection_outer_boxes'])
final_outputs['detection_outer_boxes'] = detections[
'detection_outer_boxes'
]
else:
# For RetinaNet model, apply export_config.
if isinstance(self.params.task.model, configs.retinanet.RetinaNet):
Expand All @@ -250,7 +281,7 @@ def serve(self, images: tf.Tensor):
detections = self._normalize_coordinates(detections, keys, image_info)
final_outputs = {
'decoded_boxes': detections['decoded_boxes'],
'decoded_box_scores': detections['decoded_box_scores']
'decoded_box_scores': detections['decoded_box_scores'],
}

if 'detection_masks' in detections.keys():
Expand Down
75 changes: 46 additions & 29 deletions official/vision/serving/export_saved_model_lib.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
r"""Vision models export utility function for serving/inference."""

import os
from typing import Optional, List, Union, Text, Dict
from typing import Dict, List, Optional, Union

from absl import logging
import tensorflow as tf, tf_keras
Expand Down Expand Up @@ -45,7 +45,7 @@ def export_inference_graph(
log_model_flops_and_params: bool = False,
checkpoint: Optional[tf.train.Checkpoint] = None,
input_name: Optional[str] = None,
function_keys: Optional[Union[List[Text], Dict[Text, Text]]] = None,
function_keys: Optional[Union[List[str], Dict[str, str]]] = None,
add_tpu_function_alias: Optional[bool] = False,
):
"""Exports inference graph for the model specified in the exp config.
Expand Down Expand Up @@ -83,57 +83,68 @@ def export_inference_graph(

if export_checkpoint_subdir:
output_checkpoint_directory = os.path.join(
export_dir, export_checkpoint_subdir)
export_dir, export_checkpoint_subdir
)
else:
output_checkpoint_directory = None

if export_saved_model_subdir:
output_saved_model_directory = os.path.join(
export_dir, export_saved_model_subdir)
export_dir, export_saved_model_subdir
)
else:
output_saved_model_directory = export_dir

# TODO(arashwan): Offers a direct path to use ExportModule with Task objects.
if not export_module:
if isinstance(params.task,
configs.image_classification.ImageClassificationTask):
if isinstance(
params.task, configs.image_classification.ImageClassificationTask
):
export_module = image_classification.ClassificationModule(
params=params,
batch_size=batch_size,
input_image_size=input_image_size,
input_type=input_type,
num_channels=num_channels,
input_name=input_name)
input_name=input_name,
)
elif isinstance(params.task, configs.retinanet.RetinaNetTask) or isinstance(
params.task, configs.maskrcnn.MaskRCNNTask):
params.task, configs.maskrcnn.MaskRCNNTask
):
export_module = detection.DetectionModule(
params=params,
batch_size=batch_size,
input_image_size=input_image_size,
input_type=input_type,
num_channels=num_channels,
input_name=input_name)
elif isinstance(params.task,
configs.semantic_segmentation.SemanticSegmentationTask):
input_name=input_name,
)
elif isinstance(
params.task, configs.semantic_segmentation.SemanticSegmentationTask
):
export_module = semantic_segmentation.SegmentationModule(
params=params,
batch_size=batch_size,
input_image_size=input_image_size,
input_type=input_type,
num_channels=num_channels,
input_name=input_name)
elif isinstance(params.task,
configs.video_classification.VideoClassificationTask):
input_name=input_name,
)
elif isinstance(
params.task, configs.video_classification.VideoClassificationTask
):
export_module = video_classification.VideoClassificationModule(
params=params,
batch_size=batch_size,
input_image_size=input_image_size,
input_type=input_type,
num_channels=num_channels,
input_name=input_name)
input_name=input_name,
)
else:
raise ValueError('Export module not implemented for {} task.'.format(
type(params.task)))
raise ValueError(
'Export module not implemented for {} task.'.format(type(params.task))
)

if add_tpu_function_alias:
if input_type == 'image_tensor':
Expand All @@ -160,7 +171,8 @@ def export_inference_graph(
checkpoint=checkpoint,
checkpoint_path=checkpoint_path,
timestamped=False,
save_options=save_options)
save_options=save_options,
)

if output_checkpoint_directory:
ckpt = tf.train.Checkpoint(model=export_module.model)
Expand All @@ -171,16 +183,16 @@ def export_inference_graph(
inputs_kwargs = None
if isinstance(
params.task,
(configs.retinanet.RetinaNetTask, configs.maskrcnn.MaskRCNNTask)):
(configs.retinanet.RetinaNetTask, configs.maskrcnn.MaskRCNNTask),
):
# We need to create inputs_kwargs argument to specify the input shapes for
# subclass model that overrides model.call to take multiple inputs,
# e.g., RetinaNet model.
inputs_kwargs = {
'images':
tf.TensorSpec([1] + input_image_size + [num_channels],
tf.float32),
'image_shape':
tf.TensorSpec([1, 2], tf.float32)
'images': tf.TensorSpec(
[1] + input_image_size + [num_channels], tf.float32
),
'image_shape': tf.TensorSpec([1, 2], tf.float32),
}
dummy_inputs = {
k: tf.ones(v.shape.as_list(), tf.float32)
Expand All @@ -191,9 +203,14 @@ def export_inference_graph(
else:
logging.info(
'Logging model flops and params not implemented for %s task.',
type(params.task))
type(params.task),
)
return
train_utils.try_count_flops(export_module.model, inputs_kwargs,
os.path.join(export_dir, 'model_flops.txt'))
train_utils.write_model_params(export_module.model,
os.path.join(export_dir, 'model_params.txt'))
train_utils.try_count_flops(
export_module.model,
inputs_kwargs,
os.path.join(export_dir, 'model_flops.txt'),
)
train_utils.write_model_params(
export_module.model, os.path.join(export_dir, 'model_params.txt')
)

0 comments on commit c9d1883

Please sign in to comment.