"magic_pdf/vscode:/vscode.git/clone" did not exist on "2c75a3752c1dbc24190b30845995c9985acbc945"
model_lib.py 39.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
15
r"""Constructs model, inputs, and training environment."""
16
17
18
19
20

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

21
import copy
22
import functools
23
import os
24
25
26
27

import tensorflow as tf

from object_detection import eval_util
28
from object_detection import exporter as exporter_lib
29
from object_detection import inputs
30
from object_detection.builders import graph_rewriter_builder
31
32
33
34
35
from object_detection.builders import model_builder
from object_detection.builders import optimizer_builder
from object_detection.core import standard_fields as fields
from object_detection.utils import config_util
from object_detection.utils import label_map_util
36
from object_detection.utils import ops
37
38
39
40
from object_detection.utils import shape_utils
from object_detection.utils import variables_helper
from object_detection.utils import visualization_utils as vis_utils

41
42
43
44
45
46
47
48
49
50
51
52
# pylint: disable=g-import-not-at-top
try:
  from tensorflow.contrib import framework as contrib_framework
  from tensorflow.contrib import layers as contrib_layers
  from tensorflow.contrib import learn as contrib_learn
  from tensorflow.contrib import tpu as contrib_tpu
  from tensorflow.contrib import training as contrib_training
except ImportError:
  # TF 2.0 doesn't ship with contrib.
  pass
# pylint: enable=g-import-not-at-top

53
54
55
56
57
58
59
60
# A map of names to methods that help build the model.
MODEL_BUILD_UTIL_MAP = {
    'get_configs_from_pipeline_file':
        config_util.get_configs_from_pipeline_file,
    'create_pipeline_proto_from_configs':
        config_util.create_pipeline_proto_from_configs,
    'merge_external_params_with_configs':
        config_util.merge_external_params_with_configs,
61
62
63
64
65
66
    'create_train_input_fn':
        inputs.create_train_input_fn,
    'create_eval_input_fn':
        inputs.create_eval_input_fn,
    'create_predict_input_fn':
        inputs.create_predict_input_fn,
67
    'detection_model_fn_base': model_builder.build,
68
69
70
}


71
72
def _prepare_groundtruth_for_eval(detection_model, class_agnostic,
                                  max_number_of_boxes):
73
  """Extracts groundtruth data from detection_model and prepares it for eval.
74
75
76
77

  Args:
    detection_model: A `DetectionModel` object.
    class_agnostic: Whether the detections are class_agnostic.
78
    max_number_of_boxes: Max number of groundtruth boxes.
79
80
81
82

  Returns:
    A tuple of:
    groundtruth: Dictionary with the following fields:
83
84
85
86
87
      'groundtruth_boxes': [batch_size, num_boxes, 4] float32 tensor of boxes,
        in normalized coordinates.
      'groundtruth_classes': [batch_size, num_boxes] int64 tensor of 1-indexed
        classes.
      'groundtruth_masks': 4D float32 tensor of instance masks (if provided in
88
        groundtruth)
89
90
      'groundtruth_is_crowd': [batch_size, num_boxes] bool tensor indicating
        is_crowd annotations (if provided in groundtruth).
91
92
93
      'groundtruth_area': [batch_size, num_boxes] float32 tensor indicating
        the area (in the original absolute coordinates) of annotations (if
        provided in groundtruth).
94
95
      'num_groundtruth_boxes': [batch_size] tensor containing the maximum number
        of groundtruth boxes per image..
96
97
      'groundtruth_keypoints': [batch_size, num_boxes, num_keypoints, 2] float32
        tensor of keypoints (if provided in groundtruth).
98
99
100
    class_agnostic: Boolean indicating whether detections are class agnostic.
  """
  input_data_fields = fields.InputDataFields()
101
102
103
  groundtruth_boxes = tf.stack(
      detection_model.groundtruth_lists(fields.BoxListFields.boxes))
  groundtruth_boxes_shape = tf.shape(groundtruth_boxes)
104
105
106
  # For class-agnostic models, groundtruth one-hot encodings collapse to all
  # ones.
  if class_agnostic:
107
108
    groundtruth_classes_one_hot = tf.ones(
        [groundtruth_boxes_shape[0], groundtruth_boxes_shape[1], 1])
109
  else:
110
111
    groundtruth_classes_one_hot = tf.stack(
        detection_model.groundtruth_lists(fields.BoxListFields.classes))
112
113
  label_id_offset = 1  # Applying label id offset (b/63711816)
  groundtruth_classes = (
114
      tf.argmax(groundtruth_classes_one_hot, axis=2) + label_id_offset)
115
116
117
118
119
  groundtruth = {
      input_data_fields.groundtruth_boxes: groundtruth_boxes,
      input_data_fields.groundtruth_classes: groundtruth_classes
  }
  if detection_model.groundtruth_has_field(fields.BoxListFields.masks):
120
121
122
    groundtruth[input_data_fields.groundtruth_instance_masks] = tf.stack(
        detection_model.groundtruth_lists(fields.BoxListFields.masks))

123
  if detection_model.groundtruth_has_field(fields.BoxListFields.is_crowd):
124
125
126
    groundtruth[input_data_fields.groundtruth_is_crowd] = tf.stack(
        detection_model.groundtruth_lists(fields.BoxListFields.is_crowd))

127
128
129
130
131
132
133
134
135
136
137
138
139
140
  if detection_model.groundtruth_has_field(input_data_fields.groundtruth_area):
    groundtruth[input_data_fields.groundtruth_area] = tf.stack(
        detection_model.groundtruth_lists(input_data_fields.groundtruth_area))

  if detection_model.groundtruth_has_field(fields.BoxListFields.keypoints):
    groundtruth[input_data_fields.groundtruth_keypoints] = tf.stack(
        detection_model.groundtruth_lists(fields.BoxListFields.keypoints))

  if detection_model.groundtruth_has_field(
      fields.BoxListFields.keypoint_visibilities):
    groundtruth[input_data_fields.groundtruth_keypoint_visibilities] = tf.stack(
        detection_model.groundtruth_lists(
            fields.BoxListFields.keypoint_visibilities))

141
142
  groundtruth[input_data_fields.num_groundtruth_boxes] = (
      tf.tile([max_number_of_boxes], multiples=[groundtruth_boxes_shape[0]]))
143
144
145
146
147
148
149
  return groundtruth


def unstack_batch(tensor_dict, unpad_groundtruth_tensors=True):
  """Unstacks all tensors in `tensor_dict` along 0th dimension.

  Unstacks tensor from the tensor dict along 0th dimension and returns a
150
  tensor_dict containing values that are lists of unstacked, unpadded tensors.
151
152
153
154
155
156

  Tensors in the `tensor_dict` are expected to be of one of the three shapes:
  1. [batch_size]
  2. [batch_size, height, width, channels]
  3. [batch_size, num_boxes, d1, d2, ... dn]

157
158
  When unpad_groundtruth_tensors is set to true, unstacked tensors of form 3
  above are sliced along the `num_boxes` dimension using the value in tensor
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
  field.InputDataFields.num_groundtruth_boxes.

  Note that this function has a static list of input data fields and has to be
  kept in sync with the InputDataFields defined in core/standard_fields.py

  Args:
    tensor_dict: A dictionary of batched groundtruth tensors.
    unpad_groundtruth_tensors: Whether to remove padding along `num_boxes`
      dimension of the groundtruth tensors.

  Returns:
    A dictionary where the keys are from fields.InputDataFields and values are
    a list of unstacked (optionally unpadded) tensors.

  Raises:
    ValueError: If unpad_tensors is True and `tensor_dict` does not contain
      `num_groundtruth_boxes` tensor.
  """
177
178
179
  unbatched_tensor_dict = {
      key: tf.unstack(tensor) for key, tensor in tensor_dict.items()
  }
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
  if unpad_groundtruth_tensors:
    if (fields.InputDataFields.num_groundtruth_boxes not in
        unbatched_tensor_dict):
      raise ValueError('`num_groundtruth_boxes` not found in tensor_dict. '
                       'Keys available: {}'.format(
                           unbatched_tensor_dict.keys()))
    unbatched_unpadded_tensor_dict = {}
    unpad_keys = set([
        # List of input data fields that are padded along the num_boxes
        # dimension. This list has to be kept in sync with InputDataFields in
        # standard_fields.py.
        fields.InputDataFields.groundtruth_instance_masks,
        fields.InputDataFields.groundtruth_classes,
        fields.InputDataFields.groundtruth_boxes,
        fields.InputDataFields.groundtruth_keypoints,
195
        fields.InputDataFields.groundtruth_keypoint_visibilities,
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
        fields.InputDataFields.groundtruth_group_of,
        fields.InputDataFields.groundtruth_difficult,
        fields.InputDataFields.groundtruth_is_crowd,
        fields.InputDataFields.groundtruth_area,
        fields.InputDataFields.groundtruth_weights
    ]).intersection(set(unbatched_tensor_dict.keys()))

    for key in unpad_keys:
      unpadded_tensor_list = []
      for num_gt, padded_tensor in zip(
          unbatched_tensor_dict[fields.InputDataFields.num_groundtruth_boxes],
          unbatched_tensor_dict[key]):
        tensor_shape = shape_utils.combined_static_and_dynamic_shape(
            padded_tensor)
        slice_begin = tf.zeros([len(tensor_shape)], dtype=tf.int32)
        slice_size = tf.stack(
            [num_gt] + [-1 if dim is None else dim for dim in tensor_shape[1:]])
        unpadded_tensor = tf.slice(padded_tensor, slice_begin, slice_size)
        unpadded_tensor_list.append(unpadded_tensor)
      unbatched_unpadded_tensor_dict[key] = unpadded_tensor_list
    unbatched_tensor_dict.update(unbatched_unpadded_tensor_dict)

  return unbatched_tensor_dict


pkulzc's avatar
pkulzc committed
221
def provide_groundtruth(model, labels):
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
  """Provides the labels to a model as groundtruth.

  This helper function extracts the corresponding boxes, classes,
  keypoints, weights, masks, etc. from the labels, and provides it
  as groundtruth to the models.

  Args:
    model: The detection model to provide groundtruth to.
    labels: The labels for the training or evaluation inputs.
  """
  gt_boxes_list = labels[fields.InputDataFields.groundtruth_boxes]
  gt_classes_list = labels[fields.InputDataFields.groundtruth_classes]
  gt_masks_list = None
  if fields.InputDataFields.groundtruth_instance_masks in labels:
    gt_masks_list = labels[
        fields.InputDataFields.groundtruth_instance_masks]
  gt_keypoints_list = None
  if fields.InputDataFields.groundtruth_keypoints in labels:
    gt_keypoints_list = labels[fields.InputDataFields.groundtruth_keypoints]
241
242
243
244
  gt_keypoint_visibilities_list = None
  if fields.InputDataFields.groundtruth_keypoint_visibilities in labels:
    gt_keypoint_visibilities_list = labels[
        fields.InputDataFields.groundtruth_keypoint_visibilities]
245
246
247
248
249
250
251
252
253
254
  gt_weights_list = None
  if fields.InputDataFields.groundtruth_weights in labels:
    gt_weights_list = labels[fields.InputDataFields.groundtruth_weights]
  gt_confidences_list = None
  if fields.InputDataFields.groundtruth_confidences in labels:
    gt_confidences_list = labels[
        fields.InputDataFields.groundtruth_confidences]
  gt_is_crowd_list = None
  if fields.InputDataFields.groundtruth_is_crowd in labels:
    gt_is_crowd_list = labels[fields.InputDataFields.groundtruth_is_crowd]
255
256
257
258
259
260
261
  gt_area_list = None
  if fields.InputDataFields.groundtruth_area in labels:
    gt_area_list = labels[fields.InputDataFields.groundtruth_area]
  gt_labeled_classes = None
  if fields.InputDataFields.groundtruth_labeled_classes in labels:
    gt_labeled_classes = labels[
        fields.InputDataFields.groundtruth_labeled_classes]
262
263
264
265
  model.provide_groundtruth(
      groundtruth_boxes_list=gt_boxes_list,
      groundtruth_classes_list=gt_classes_list,
      groundtruth_confidences_list=gt_confidences_list,
266
      groundtruth_labeled_classes=gt_labeled_classes,
267
268
      groundtruth_masks_list=gt_masks_list,
      groundtruth_keypoints_list=gt_keypoints_list,
269
      groundtruth_keypoint_visibilities_list=gt_keypoint_visibilities_list,
270
      groundtruth_weights_list=gt_weights_list,
271
272
      groundtruth_is_crowd_list=gt_is_crowd_list,
      groundtruth_area_list=gt_area_list)
273
274


275
276
def create_model_fn(detection_model_fn, configs, hparams, use_tpu=False,
                    postprocess_on_cpu=False):
277
278
279
280
281
282
283
284
  """Creates a model function for `Estimator`.

  Args:
    detection_model_fn: Function that returns a `DetectionModel` instance.
    configs: Dictionary of pipeline config objects.
    hparams: `HParams` object.
    use_tpu: Boolean indicating whether model should be constructed for
        use on TPU.
285
286
    postprocess_on_cpu: When use_tpu and postprocess_on_cpu is true, postprocess
        is scheduled on the host cpu.
287
288
289
290
291
292

  Returns:
    `model_fn` for `Estimator`.
  """
  train_config = configs['train_config']
  eval_input_config = configs['eval_input_config']
293
  eval_config = configs['eval_config']
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311

  def model_fn(features, labels, mode, params=None):
    """Constructs the object detection model.

    Args:
      features: Dictionary of feature tensors, returned from `input_fn`.
      labels: Dictionary of groundtruth tensors if mode is TRAIN or EVAL,
        otherwise None.
      mode: Mode key from tf.estimator.ModeKeys.
      params: Parameter dictionary passed from the estimator.

    Returns:
      An `EstimatorSpec` that encapsulates the model and its serving
        configurations.
    """
    params = params or {}
    total_loss, train_op, detections, export_outputs = None, None, None, None
    is_training = mode == tf.estimator.ModeKeys.TRAIN
312
313
314
315

    # Make sure to set the Keras learning phase. True during training,
    # False for inference.
    tf.keras.backend.set_learning_phase(is_training)
316
317
318
319
320
321
322
    # Set policy for mixed-precision training with Keras-based models.
    if use_tpu and train_config.use_bfloat16:
      from tensorflow.python.keras.engine import base_layer_utils  # pylint: disable=g-import-not-at-top
      # Enable v2 behavior, as `mixed_bfloat16` is only supported in TF 2.0.
      base_layer_utils.enable_v2_dtype_behavior()
      tf.compat.v2.keras.mixed_precision.experimental.set_policy(
          'mixed_bfloat16')
323
324
    detection_model = detection_model_fn(
        is_training=is_training, add_summaries=(not use_tpu))
325
326
327
328
329
330
331
    scaffold_fn = None

    if mode == tf.estimator.ModeKeys.TRAIN:
      labels = unstack_batch(
          labels,
          unpad_groundtruth_tensors=train_config.unpad_groundtruth_tensors)
    elif mode == tf.estimator.ModeKeys.EVAL:
332
333
334
335
336
      # For evaling on train data, it is necessary to check whether groundtruth
      # must be unpadded.
      boxes_shape = (
          labels[fields.InputDataFields.groundtruth_boxes].get_shape()
          .as_list())
337
      unpad_groundtruth_tensors = boxes_shape[1] is not None and not use_tpu
338
339
      labels = unstack_batch(
          labels, unpad_groundtruth_tensors=unpad_groundtruth_tensors)
340
341

    if mode in (tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL):
pkulzc's avatar
pkulzc committed
342
      provide_groundtruth(detection_model, labels)
343
344

    preprocessed_images = features[fields.InputDataFields.image]
345
346
347

    side_inputs = detection_model.get_side_inputs(features)

348
    if use_tpu and train_config.use_bfloat16:
349
      with contrib_tpu.bfloat16_scope():
350
351
        prediction_dict = detection_model.predict(
            preprocessed_images,
352
            features[fields.InputDataFields.true_image_shape], **side_inputs)
353
        prediction_dict = ops.bfloat16_to_float32_nested(prediction_dict)
354
355
356
    else:
      prediction_dict = detection_model.predict(
          preprocessed_images,
357
          features[fields.InputDataFields.true_image_shape], **side_inputs)
358
359
360
361

    def postprocess_wrapper(args):
      return detection_model.postprocess(args[0], args[1])

362
    if mode in (tf.estimator.ModeKeys.EVAL, tf.estimator.ModeKeys.PREDICT):
363
      if use_tpu and postprocess_on_cpu:
364
        detections = contrib_tpu.outside_compilation(
365
366
367
368
369
370
371
            postprocess_wrapper,
            (prediction_dict,
             features[fields.InputDataFields.true_image_shape]))
      else:
        detections = postprocess_wrapper((
            prediction_dict,
            features[fields.InputDataFields.true_image_shape]))
372
373

    if mode == tf.estimator.ModeKeys.TRAIN:
374
375
      load_pretrained = hparams.load_pretrained if hparams else False
      if train_config.fine_tune_checkpoint and load_pretrained:
376
377
378
379
380
381
382
383
        if not train_config.fine_tune_checkpoint_type:
          # train_config.from_detection_checkpoint field is deprecated. For
          # backward compatibility, set train_config.fine_tune_checkpoint_type
          # based on train_config.from_detection_checkpoint.
          if train_config.from_detection_checkpoint:
            train_config.fine_tune_checkpoint_type = 'detection'
          else:
            train_config.fine_tune_checkpoint_type = 'classification'
384
        asg_map = detection_model.restore_map(
385
            fine_tune_checkpoint_type=train_config.fine_tune_checkpoint_type,
386
387
388
389
            load_all_detection_checkpoint_vars=(
                train_config.load_all_detection_checkpoint_vars))
        available_var_map = (
            variables_helper.get_variables_available_in_checkpoint(
390
391
                asg_map,
                train_config.fine_tune_checkpoint,
392
393
                include_global_step=False))
        if use_tpu:
394

395
396
397
398
          def tpu_scaffold():
            tf.train.init_from_checkpoint(train_config.fine_tune_checkpoint,
                                          available_var_map)
            return tf.train.Scaffold()
399

400
401
402
403
404
405
          scaffold_fn = tpu_scaffold
        else:
          tf.train.init_from_checkpoint(train_config.fine_tune_checkpoint,
                                        available_var_map)

    if mode in (tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL):
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
      if (mode == tf.estimator.ModeKeys.EVAL and
          eval_config.use_dummy_loss_in_eval):
        total_loss = tf.constant(1.0)
        losses_dict = {'Loss/total_loss': total_loss}
      else:
        losses_dict = detection_model.loss(
            prediction_dict, features[fields.InputDataFields.true_image_shape])
        losses = [loss_tensor for loss_tensor in losses_dict.values()]
        if train_config.add_regularization_loss:
          regularization_losses = detection_model.regularization_losses()
          if use_tpu and train_config.use_bfloat16:
            regularization_losses = ops.bfloat16_to_float32_nested(
                regularization_losses)
          if regularization_losses:
            regularization_loss = tf.add_n(
                regularization_losses, name='regularization_loss')
            losses.append(regularization_loss)
            losses_dict['Loss/regularization_loss'] = regularization_loss
        total_loss = tf.add_n(losses, name='total_loss')
        losses_dict['Loss/total_loss'] = total_loss
426

427
428
429
430
431
      if 'graph_rewriter_config' in configs:
        graph_rewriter_fn = graph_rewriter_builder.build(
            configs['graph_rewriter_config'], is_training=is_training)
        graph_rewriter_fn()

432
433
      # TODO(rathodv): Stop creating optimizer summary vars in EVAL mode once we
      # can write learning rate summaries on TPU without host calls.
434
435
436
437
      global_step = tf.train.get_or_create_global_step()
      training_optimizer, optimizer_summary_vars = optimizer_builder.build(
          train_config.optimizer)

438
    if mode == tf.estimator.ModeKeys.TRAIN:
439
      if use_tpu:
440
        training_optimizer = contrib_tpu.CrossShardOptimizer(training_optimizer)
441
442
443

      # Optionally freeze some layers by setting their gradients to be zero.
      trainable_variables = None
444
445
446
447
448
449
      include_variables = (
          train_config.update_trainable_variables
          if train_config.update_trainable_variables else None)
      exclude_variables = (
          train_config.freeze_variables
          if train_config.freeze_variables else None)
450
      trainable_variables = contrib_framework.filter_variables(
451
452
453
          tf.trainable_variables(),
          include_patterns=include_variables,
          exclude_patterns=exclude_variables)
454
455
456
457
458
459
460
461
462

      clip_gradients_value = None
      if train_config.gradient_clipping_by_norm > 0:
        clip_gradients_value = train_config.gradient_clipping_by_norm

      if not use_tpu:
        for var in optimizer_summary_vars:
          tf.summary.scalar(var.op.name, var)
      summaries = [] if use_tpu else None
463
464
      if train_config.summarize_gradients:
        summaries = ['gradients', 'gradient_norm', 'global_gradient_norm']
465
      train_op = contrib_layers.optimize_loss(
466
467
468
469
470
          loss=total_loss,
          global_step=global_step,
          learning_rate=None,
          clip_gradients=clip_gradients_value,
          optimizer=training_optimizer,
471
          update_ops=detection_model.updates(),
472
473
474
475
476
          variables=trainable_variables,
          summaries=summaries,
          name='')  # Preventing scope prefix on all variables.

    if mode == tf.estimator.ModeKeys.PREDICT:
477
      exported_output = exporter_lib.add_output_tensor_nodes(detections)
478
479
      export_outputs = {
          tf.saved_model.signature_constants.PREDICT_METHOD_NAME:
480
              tf.estimator.export.PredictOutput(exported_output)
481
482
483
      }

    eval_metric_ops = None
484
    scaffold = None
485
    if mode == tf.estimator.ModeKeys.EVAL:
486
487
      class_agnostic = (
          fields.DetectionResultFields.detection_classes not in detections)
488
489
490
      groundtruth = _prepare_groundtruth_for_eval(
          detection_model, class_agnostic,
          eval_input_config.max_number_of_boxes)
491
      use_original_images = fields.InputDataFields.original_image in features
pkulzc's avatar
pkulzc committed
492
      if use_original_images:
493
494
495
496
497
        eval_images = features[fields.InputDataFields.original_image]
        true_image_shapes = tf.slice(
            features[fields.InputDataFields.true_image_shape], [0, 0], [-1, 3])
        original_image_spatial_shapes = features[fields.InputDataFields
                                                 .original_image_spatial_shape]
pkulzc's avatar
pkulzc committed
498
499
      else:
        eval_images = features[fields.InputDataFields.image]
500
501
        true_image_shapes = None
        original_image_spatial_shapes = None
pkulzc's avatar
pkulzc committed
502

503
504
505
      eval_dict = eval_util.result_dict_for_batched_example(
          eval_images,
          features[inputs.HASH_KEY],
506
507
508
          detections,
          groundtruth,
          class_agnostic=class_agnostic,
509
510
511
          scale_to_absolute=True,
          original_image_spatial_shapes=original_image_spatial_shapes,
          true_image_shapes=true_image_shapes)
512

513
514
515
516
      if fields.InputDataFields.image_additional_channels in features:
        eval_dict[fields.InputDataFields.image_additional_channels] = features[
            fields.InputDataFields.image_additional_channels]

517
518
519
520
521
      if class_agnostic:
        category_index = label_map_util.create_class_agnostic_category_index()
      else:
        category_index = label_map_util.create_category_index_from_labelmap(
            eval_input_config.label_map_path)
522
      vis_metric_ops = None
523
      if not use_tpu and use_original_images:
524
525
526
        keypoint_edges = [
            (kp.start, kp.end) for kp in eval_config.keypoint_edge]

527
528
529
530
531
        eval_metric_op_vis = vis_utils.VisualizeSingleFrameDetections(
            category_index,
            max_examples_to_draw=eval_config.num_visualizations,
            max_boxes_to_draw=eval_config.max_num_boxes_to_visualize,
            min_score_thresh=eval_config.min_score_threshold,
532
533
            use_normalized_coordinates=False,
            keypoint_edges=keypoint_edges or None)
534
535
        vis_metric_ops = eval_metric_op_vis.get_estimator_eval_metric_ops(
            eval_dict)
536

537
538
      # Eval metrics on a single example.
      eval_metric_ops = eval_util.get_eval_metric_ops_for_evaluators(
DefineFC's avatar
DefineFC committed
539
          eval_config, list(category_index.values()), eval_dict)
540
541
542
543
      for loss_key, loss_tensor in iter(losses_dict.items()):
        eval_metric_ops[loss_key] = tf.metrics.mean(loss_tensor)
      for var in optimizer_summary_vars:
        eval_metric_ops[var.op.name] = (var, tf.no_op())
544
545
      if vis_metric_ops is not None:
        eval_metric_ops.update(vis_metric_ops)
546
      eval_metric_ops = {str(k): v for k, v in eval_metric_ops.items()}
547

548
549
550
551
552
553
554
555
556
557
      if eval_config.use_moving_averages:
        variable_averages = tf.train.ExponentialMovingAverage(0.0)
        variables_to_restore = variable_averages.variables_to_restore()
        keep_checkpoint_every_n_hours = (
            train_config.keep_checkpoint_every_n_hours)
        saver = tf.train.Saver(
            variables_to_restore,
            keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours)
        scaffold = tf.train.Scaffold(saver=saver)

558
559
    # EVAL executes on CPU, so use regular non-TPU EstimatorSpec.
    if use_tpu and mode != tf.estimator.ModeKeys.EVAL:
560
      return contrib_tpu.TPUEstimatorSpec(
561
562
563
564
565
566
567
568
          mode=mode,
          scaffold_fn=scaffold_fn,
          predictions=detections,
          loss=total_loss,
          train_op=train_op,
          eval_metrics=eval_metric_ops,
          export_outputs=export_outputs)
    else:
569
570
571
572
573
574
575
576
577
      if scaffold is None:
        keep_checkpoint_every_n_hours = (
            train_config.keep_checkpoint_every_n_hours)
        saver = tf.train.Saver(
            sharded=True,
            keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours,
            save_relative_paths=True)
        tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
        scaffold = tf.train.Scaffold(saver=saver)
578
579
580
581
582
583
      return tf.estimator.EstimatorSpec(
          mode=mode,
          predictions=detections,
          loss=total_loss,
          train_op=train_op,
          eval_metric_ops=eval_metric_ops,
584
585
          export_outputs=export_outputs,
          scaffold=scaffold)
586
587
588
589

  return model_fn


590
591
592
def create_estimator_and_inputs(run_config,
                                hparams,
                                pipeline_config_path,
593
                                config_override=None,
594
                                train_steps=None,
595
                                sample_1_of_n_eval_examples=1,
596
                                sample_1_of_n_eval_on_train_examples=1,
597
598
599
600
601
                                model_fn_creator=create_model_fn,
                                use_tpu_estimator=False,
                                use_tpu=False,
                                num_shards=1,
                                params=None,
602
                                override_eval_num_epochs=True,
603
                                save_final_config=False,
604
605
                                postprocess_on_cpu=False,
                                export_to_tpu=None,
606
607
                                **kwargs):
  """Creates `Estimator`, input functions, and steps.
608
609
610
611
612

  Args:
    run_config: A `RunConfig`.
    hparams: A `HParams`.
    pipeline_config_path: A path to a pipeline config file.
613
614
    config_override: A pipeline_pb2.TrainEvalPipelineConfig text proto to
      override the config from `pipeline_config_path`.
615
616
    train_steps: Number of training steps. If None, the number of training steps
      is set from the `TrainConfig` proto.
617
618
619
620
621
    sample_1_of_n_eval_examples: Integer representing how often an eval example
      should be sampled. If 1, will sample all examples.
    sample_1_of_n_eval_on_train_examples: Similar to
      `sample_1_of_n_eval_examples`, except controls the sampling of training
      data for evaluation.
622
623
624
625
626
627
628
629
630
631
    model_fn_creator: A function that creates a `model_fn` for `Estimator`.
      Follows the signature:

      * Args:
        * `detection_model_fn`: Function that returns `DetectionModel` instance.
        * `configs`: Dictionary of pipeline config objects.
        * `hparams`: `HParams` object.
      * Returns:
        `model_fn` for `Estimator`.

632
633
634
635
636
637
638
639
    use_tpu_estimator: Whether a `TPUEstimator` should be returned. If False,
      an `Estimator` will be returned.
    use_tpu: Boolean, whether training and evaluation should run on TPU. Only
      used if `use_tpu_estimator` is True.
    num_shards: Number of shards (TPU cores). Only used if `use_tpu_estimator`
      is True.
    params: Parameter dictionary passed from the estimator. Only used if
      `use_tpu_estimator` is True.
640
641
    override_eval_num_epochs: Whether to overwrite the number of epochs to 1 for
      eval_input.
642
643
    save_final_config: Whether to save final config (obtained after applying
      overrides) to `estimator.model_dir`.
644
645
646
647
648
    postprocess_on_cpu: When use_tpu and postprocess_on_cpu are true,
      postprocess is scheduled on the host cpu.
    export_to_tpu: When use_tpu and export_to_tpu are true,
      `export_savedmodel()` exports a metagraph for serving on TPU besides the
      one on CPU.
649
650
651
    **kwargs: Additional keyword arguments for configuration override.

  Returns:
652
653
654
    A dictionary with the following fields:
    'estimator': An `Estimator` or `TPUEstimator`.
    'train_input_fn': A training input function.
655
656
    'eval_input_fns': A list of all evaluation input functions.
    'eval_input_names': A list of names for each evaluation input.
657
    'eval_on_train_input_fn': An evaluation-on-train input function.
658
659
660
    'predict_input_fn': A prediction input function.
    'train_steps': Number of training steps. Either directly from input or from
      configuration.
661
  """
662
663
664
665
  get_configs_from_pipeline_file = MODEL_BUILD_UTIL_MAP[
      'get_configs_from_pipeline_file']
  merge_external_params_with_configs = MODEL_BUILD_UTIL_MAP[
      'merge_external_params_with_configs']
666
667
  create_pipeline_proto_from_configs = MODEL_BUILD_UTIL_MAP[
      'create_pipeline_proto_from_configs']
668
669
670
  create_train_input_fn = MODEL_BUILD_UTIL_MAP['create_train_input_fn']
  create_eval_input_fn = MODEL_BUILD_UTIL_MAP['create_eval_input_fn']
  create_predict_input_fn = MODEL_BUILD_UTIL_MAP['create_predict_input_fn']
671
  detection_model_fn_base = MODEL_BUILD_UTIL_MAP['detection_model_fn_base']
672

673
674
  configs = get_configs_from_pipeline_file(
      pipeline_config_path, config_override=config_override)
675
676
  kwargs.update({
      'train_steps': train_steps,
677
      'use_bfloat16': configs['train_config'].use_bfloat16 and use_tpu
678
  })
pkulzc's avatar
pkulzc committed
679
680
681
682
  if sample_1_of_n_eval_examples >= 1:
    kwargs.update({
        'sample_1_of_n_eval_examples': sample_1_of_n_eval_examples
    })
683
684
685
686
  if override_eval_num_epochs:
    kwargs.update({'eval_num_epochs': 1})
    tf.logging.warning(
        'Forced number of epochs for all eval validations to be 1.')
687
  configs = merge_external_params_with_configs(
688
      configs, hparams, kwargs_dict=kwargs)
689
690
691
692
  model_config = configs['model']
  train_config = configs['train_config']
  train_input_config = configs['train_input_config']
  eval_config = configs['eval_config']
693
694
695
696
697
698
699
700
701
702
703
  eval_input_configs = configs['eval_input_configs']
  eval_on_train_input_config = copy.deepcopy(train_input_config)
  eval_on_train_input_config.sample_1_of_n_examples = (
      sample_1_of_n_eval_on_train_examples)
  if override_eval_num_epochs and eval_on_train_input_config.num_epochs != 1:
    tf.logging.warning('Expected number of evaluation epochs is 1, but '
                       'instead encountered `eval_on_train_input_config'
                       '.num_epochs` = '
                       '{}. Overwriting `num_epochs` to 1.'.format(
                           eval_on_train_input_config.num_epochs))
    eval_on_train_input_config.num_epochs = 1
704

705
706
707
  # update train_steps from config but only when non-zero value is provided
  if train_steps is None and train_config.num_steps != 0:
    train_steps = train_config.num_steps
708
709

  detection_model_fn = functools.partial(
710
      detection_model_fn_base, model_config=model_config)
711

712
  # Create the input functions for TRAIN/EVAL/PREDICT.
713
  train_input_fn = create_train_input_fn(
714
715
716
      train_config=train_config,
      train_input_config=train_input_config,
      model_config=model_config)
717
718
719
720
721
722
723
724
725
  eval_input_fns = [
      create_eval_input_fn(
          eval_config=eval_config,
          eval_input_config=eval_input_config,
          model_config=model_config) for eval_input_config in eval_input_configs
  ]
  eval_input_names = [
      eval_input_config.name for eval_input_config in eval_input_configs
  ]
726
727
  eval_on_train_input_fn = create_eval_input_fn(
      eval_config=eval_config,
728
      eval_input_config=eval_on_train_input_config,
729
      model_config=model_config)
730
  predict_input_fn = create_predict_input_fn(
731
      model_config=model_config, predict_input_config=eval_input_configs[0])
732

733
734
735
  # Read export_to_tpu from hparams if not passed.
  if export_to_tpu is None:
    export_to_tpu = hparams.get('export_to_tpu', False)
736
737
  tf.logging.info('create_estimator_and_inputs: use_tpu %s, export_to_tpu %s',
                  use_tpu, export_to_tpu)
738
739
  model_fn = model_fn_creator(detection_model_fn, configs, hparams, use_tpu,
                              postprocess_on_cpu)
740
  if use_tpu_estimator:
741
    estimator = contrib_tpu.TPUEstimator(
742
743
744
745
746
747
        model_fn=model_fn,
        train_batch_size=train_config.batch_size,
        # For each core, only batch size 1 is supported for eval.
        eval_batch_size=num_shards * 1 if use_tpu else 1,
        use_tpu=use_tpu,
        config=run_config,
748
749
        export_to_tpu=export_to_tpu,
        eval_on_tpu=False,  # Eval runs on CPU, so disable eval on TPU
pkulzc's avatar
pkulzc committed
750
        params=params if params else {})
751
752
  else:
    estimator = tf.estimator.Estimator(model_fn=model_fn, config=run_config)
753

754
  # Write the as-run pipeline config to disk.
755
  if run_config.is_chief and save_final_config:
756
    pipeline_config_final = create_pipeline_proto_from_configs(configs)
757
    config_util.save_pipeline_config(pipeline_config_final, estimator.model_dir)
758

759
  return dict(
760
761
      estimator=estimator,
      train_input_fn=train_input_fn,
762
763
      eval_input_fns=eval_input_fns,
      eval_input_names=eval_input_names,
764
      eval_on_train_input_fn=eval_on_train_input_fn,
765
      predict_input_fn=predict_input_fn,
766
      train_steps=train_steps)
767
768
769


def create_train_and_eval_specs(train_input_fn,
770
                                eval_input_fns,
771
                                eval_on_train_input_fn,
772
773
774
775
                                predict_input_fn,
                                train_steps,
                                eval_on_train_data=False,
                                final_exporter_name='Servo',
776
                                eval_spec_names=None):
777
778
779
780
  """Creates a `TrainSpec` and `EvalSpec`s.

  Args:
    train_input_fn: Function that produces features and labels on train data.
781
782
    eval_input_fns: A list of functions that produce features and labels on eval
      data.
783
784
    eval_on_train_input_fn: Function that produces features and labels for
      evaluation on train data.
785
786
787
788
789
    predict_input_fn: Function that produces features for inference.
    train_steps: Number of training steps.
    eval_on_train_data: Whether to evaluate model on training data. Default is
      False.
    final_exporter_name: String name given to `FinalExporter`.
790
    eval_spec_names: A list of string names for each `EvalSpec`.
791
792

  Returns:
793
794
795
    Tuple of `TrainSpec` and list of `EvalSpecs`. If `eval_on_train_data` is
    True, the last `EvalSpec` in the list will correspond to training data. The
    rest EvalSpecs in the list are evaluation datas.
796
797
798
799
  """
  train_spec = tf.estimator.TrainSpec(
      input_fn=train_input_fn, max_steps=train_steps)

800
  if eval_spec_names is None:
801
    eval_spec_names = [str(i) for i in range(len(eval_input_fns))]
802
803

  eval_specs = []
804
805
806
807
808
809
810
811
  for index, (eval_spec_name, eval_input_fn) in enumerate(
      zip(eval_spec_names, eval_input_fns)):
    # Uses final_exporter_name as exporter_name for the first eval spec for
    # backward compatibility.
    if index == 0:
      exporter_name = final_exporter_name
    else:
      exporter_name = '{}_{}'.format(final_exporter_name, eval_spec_name)
812
813
814
815
816
817
818
819
    exporter = tf.estimator.FinalExporter(
        name=exporter_name, serving_input_receiver_fn=predict_input_fn)
    eval_specs.append(
        tf.estimator.EvalSpec(
            name=eval_spec_name,
            input_fn=eval_input_fn,
            steps=None,
            exporters=exporter))
820
821
822
823

  if eval_on_train_data:
    eval_specs.append(
        tf.estimator.EvalSpec(
824
            name='eval_on_train', input_fn=eval_on_train_input_fn, steps=None))
825
826

  return train_spec, eval_specs
827
828


829
def continuous_eval(estimator, model_dir, input_fn, train_steps, name):
830
831
832
833
834
835
836
837
838
839
  """Perform continuous evaluation on checkpoints written to a model directory.

  Args:
    estimator: Estimator object to use for evaluation.
    model_dir: Model directory to read checkpoints for continuous evaluation.
    input_fn: Input function to use for evaluation.
    train_steps: Number of training steps. This is used to infer the last
      checkpoint and stop evaluation loop.
    name: Namescope for eval summary.
  """
840

841
842
843
844
  def terminate_eval():
    tf.logging.info('Terminating eval after 180 seconds of no checkpoints')
    return True

845
  for ckpt in contrib_training.checkpoints_iterator(
846
847
848
849
850
851
      model_dir, min_interval_secs=180, timeout=None,
      timeout_fn=terminate_eval):

    tf.logging.info('Starting Evaluation.')
    try:
      eval_results = estimator.evaluate(
852
          input_fn=input_fn, steps=None, checkpoint_path=ckpt, name=name)
853
854
855
856
857
858
859
860
861
862
863
864
865
866
      tf.logging.info('Eval results: %s' % eval_results)

      # Terminate eval job when final checkpoint is reached
      current_step = int(os.path.basename(ckpt).split('-')[1])
      if current_step >= train_steps:
        tf.logging.info(
            'Evaluation finished after training step %d' % current_step)
        break

    except tf.errors.NotFoundError:
      tf.logging.info(
          'Checkpoint %s no longer exists, skipping checkpoint' % ckpt)


867
868
869
870
871
872
873
874
def populate_experiment(run_config,
                        hparams,
                        pipeline_config_path,
                        train_steps=None,
                        eval_steps=None,
                        model_fn_creator=create_model_fn,
                        **kwargs):
  """Populates an `Experiment` object.
875

876
877
  EXPERIMENT CLASS IS DEPRECATED. Please switch to
  tf.estimator.train_and_evaluate. As an example, see model_main.py.
878

879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
  Args:
    run_config: A `RunConfig`.
    hparams: A `HParams`.
    pipeline_config_path: A path to a pipeline config file.
    train_steps: Number of training steps. If None, the number of training steps
      is set from the `TrainConfig` proto.
    eval_steps: Number of evaluation steps per evaluation cycle. If None, the
      number of evaluation steps is set from the `EvalConfig` proto.
    model_fn_creator: A function that creates a `model_fn` for `Estimator`.
      Follows the signature:

      * Args:
        * `detection_model_fn`: Function that returns `DetectionModel` instance.
        * `configs`: Dictionary of pipeline config objects.
        * `hparams`: `HParams` object.
      * Returns:
        `model_fn` for `Estimator`.

    **kwargs: Additional keyword arguments for configuration override.

  Returns:
    An `Experiment` that defines all aspects of training, evaluation, and
    export.
  """
  tf.logging.warning('Experiment is being deprecated. Please use '
                     'tf.estimator.train_and_evaluate(). See model_main.py for '
                     'an example.')
  train_and_eval_dict = create_estimator_and_inputs(
      run_config,
      hparams,
      pipeline_config_path,
      train_steps=train_steps,
      eval_steps=eval_steps,
      model_fn_creator=model_fn_creator,
913
      save_final_config=True,
914
915
916
      **kwargs)
  estimator = train_and_eval_dict['estimator']
  train_input_fn = train_and_eval_dict['train_input_fn']
917
  eval_input_fns = train_and_eval_dict['eval_input_fns']
918
919
920
921
  predict_input_fn = train_and_eval_dict['predict_input_fn']
  train_steps = train_and_eval_dict['train_steps']

  export_strategies = [
922
      contrib_learn.utils.saved_model_export_utils.make_export_strategy(
923
924
925
          serving_input_fn=predict_input_fn)
  ]

926
  return contrib_learn.Experiment(
927
928
      estimator=estimator,
      train_input_fn=train_input_fn,
929
      eval_input_fn=eval_input_fns[0],
930
      train_steps=train_steps,
931
      eval_steps=None,
932
      export_strategies=export_strategies,
933
934
      eval_delay_secs=120,
  )