Python training.create_global_step函数代码示例

  def predict(self, input_fn, predict_keys=None, hooks=None, checkpoint_path=None):
    """Returns predictions for given features.

    Args:
      input_fn: Input function returning features which is a dictionary of
        string feature name to `Tensor` or `SparseTensor`. If it returns a
        tuple, first item is extracted as features. Prediction continues until
        `input_fn` raises an end-of-input exception (`OutOfRangeError` or
        `StopIteration`).
      predict_keys: list of `str`, name of the keys to predict. It is used if
        the `EstimatorSpec.predictions` is a `dict`. If `predict_keys` is used
        then rest of the predictions will be filtered from the dictionary. If
        `None`, returns all.
      hooks: List of `SessionRunHook` subclass instances. Used for callbacks
        inside the prediction call.
      checkpoint_path: Path of a specific checkpoint to predict. If `None`, the
        latest checkpoint in `model_dir` is used.

    Yields:
      Evaluated values of `predictions` tensors.

    Raises:
      ValueError: Could not find a trained model in model_dir.
      ValueError: if batch length of predictions are not same.
      ValueError: If there is a conflict between `predict_keys` and
        `predictions`. For example if `predict_keys` is not `None` but
        `EstimatorSpec.predictions` is not a `dict`.
    """
    hooks = _check_hooks_type(hooks)
    # Check that model has been trained.
    if not checkpoint_path:
      checkpoint_path = saver.latest_checkpoint(self._model_dir)
    if not checkpoint_path:
      raise ValueError('Could not find trained model in model_dir: {}.'.format(
          self._model_dir))

    with ops.Graph().as_default() as g:
      random_seed.set_random_seed(self._config.tf_random_seed)
      training.create_global_step(g)
      features = self._get_features_from_input_fn(input_fn)
      estimator_spec = self._call_model_fn(features, None,
                                           model_fn_lib.ModeKeys.PREDICT)
      predictions = self._extract_keys(estimator_spec.predictions, predict_keys)
      with training.MonitoredSession(
          session_creator=training.ChiefSessionCreator(
              checkpoint_filename_with_path=checkpoint_path,
              scaffold=estimator_spec.scaffold,
              config=config_pb2.ConfigProto(allow_soft_placement=True)),
          hooks=hooks) as mon_sess:
        while not mon_sess.should_stop():
          preds_evaluated = mon_sess.run(predictions)
          if not isinstance(predictions, dict):
            for pred in preds_evaluated:
              yield pred
          else:
            for i in range(self._extract_batch_length(preds_evaluated)):
              yield {
                  key: value[i]
                  for key, value in six.iteritems(preds_evaluated)
              }

    def test_build_all_summaries(self):
        training.create_global_step()
        x = {'x': tf.placeholder(tf.float32, [2, 89])}
        y = tf.constant([[1], [1]])

        model = BaseModel(plx.Modes.TRAIN, graph_fn=self.get_dummy_graph_fn(),
                          loss_config=LossConfig(module='log_loss'),
                          optimizer_config=OptimizerConfig(module='adadelta',
                                                           decay_type='exponential_decay'),
                          model_type=BaseModel.Types.CLASSIFIER, eval_metrics_config=[],
                          summaries='all', name='test')

        model(x, y, None, None)

        # Only var are created
        learning_rate_summaries = 0
        activations_summaries = 0
        gradients_summaries = 0
        loss_summaries = 0

        for s_name in get_tracked(collection=tf.GraphKeys.SUMMARIES_BY_NAMES).keys():
            if 'learning_rate' in s_name:
                learning_rate_summaries += 1
            elif 'Activation' in s_name:
                activations_summaries += 1
            elif 'Loss' in s_name:
                loss_summaries += 1
            elif 'Gradient' in s_name:
                gradients_summaries += 1

        assert learning_rate_summaries > 0
        assert activations_summaries > 0
        assert gradients_summaries > 0
        assert loss_summaries > 0

  def _train_model(self, input_fn, hooks):
    all_hooks = []
    with ops.Graph().as_default() as g, g.device(self._device_fn):
      random_seed.set_random_seed(self._config.tf_random_seed)
      global_step_tensor = training.create_global_step(g)
      with ops.device('/cpu:0'):
        features, labels = input_fn()
      estimator_spec = self._call_model_fn(features, labels,
                                           model_fn_lib.ModeKeys.TRAIN)
      ops.add_to_collection(ops.GraphKeys.LOSSES, estimator_spec.loss)
      all_hooks.extend([
          training.NanTensorHook(estimator_spec.loss),
          training.LoggingTensorHook(
              {
                  'loss': estimator_spec.loss,
                  'step': global_step_tensor
              },
              every_n_iter=100)
      ])
      all_hooks.extend(hooks)
      all_hooks.extend(estimator_spec.training_hooks)

      if not (estimator_spec.scaffold.saver or
              ops.get_collection(ops.GraphKeys.SAVERS)):
        ops.add_to_collection(ops.GraphKeys.SAVERS,
                              training.Saver(
                                  sharded=True,
                                  max_to_keep=self._config.keep_checkpoint_max,
                                  defer_build=True))

      chief_hooks = []
      if (self._config.save_checkpoints_secs or
          self._config.save_checkpoints_steps):
        saver_hook_exists = any([
            isinstance(h, training.CheckpointSaverHook)
            for h in (all_hooks + chief_hooks +
                      estimator_spec.training_chief_hooks)
        ])
        if not saver_hook_exists:
          chief_hooks = [
              training.CheckpointSaverHook(
                  self._model_dir,
                  save_secs=self._config.save_checkpoints_secs,
                  save_steps=self._config.save_checkpoints_steps,
                  scaffold=estimator_spec.scaffold)
          ]
      with training.MonitoredTrainingSession(
          master=self._config.master,
          is_chief=self._config.is_chief,
          checkpoint_dir=self._model_dir,
          scaffold=estimator_spec.scaffold,
          hooks=all_hooks,
          chief_only_hooks=chief_hooks + estimator_spec.training_chief_hooks,
          save_checkpoint_secs=0,  # Saving is handled by a hook.
          save_summaries_steps=self._config.save_summary_steps,
          config=config_pb2.ConfigProto(allow_soft_placement=True)) as mon_sess:
        loss = None
        while not mon_sess.should_stop():
          _, loss = mon_sess.run([estimator_spec.train_op, estimator_spec.loss])
      return loss

  def _evaluate_model(self,
                      input_fn,
                      hooks=None,
                      checkpoint_path=None,
                      name=''):
    """Evaluates the model using the training.evaluation library."""
    # Check that model has been trained (if nothing has been set explicitly).
    if not checkpoint_path:
      latest_path = saver.latest_checkpoint(self._model_dir)
      if not latest_path:
        raise ValueError('Could not find trained model in model_dir: {}.'.
                         format(self._model_dir))
      checkpoint_path = latest_path

    # Setup output directory.
    eval_dir = os.path.join(self._model_dir, 'eval' if not name else
                            'eval_' + name)

    with ops.Graph().as_default() as g:
      random_seed.set_random_seed(self._config.tf_random_seed)
      global_step_tensor = training.create_global_step(g)
      features, labels = input_fn()
      estimator_spec = self._call_model_fn(
          features, labels, model_fn_lib.ModeKeys.EVAL)

      if model_fn_lib.MetricKeys.LOSS in estimator_spec.eval_metric_ops:
        raise ValueError(
            'Metric with name "%s" is not allowed, because Estimator ' % (
                model_fn_lib.MetricKeys.LOSS) +
            'already defines a default metric with the same name.')
      estimator_spec.eval_metric_ops[
          model_fn_lib.MetricKeys.LOSS] = metrics_lib.mean(estimator_spec.loss)

      update_op, eval_dict = _extract_metric_update_ops(
          estimator_spec.eval_metric_ops)

      if ops.GraphKeys.GLOBAL_STEP in eval_dict:
        raise ValueError(
            'Metric with name `global_step` is not allowed, because Estimator '
            'already defines a default metric with the same name.')
      eval_dict[ops.GraphKeys.GLOBAL_STEP] = global_step_tensor

      eval_results = evaluation._evaluate_once(  # pylint: disable=protected-access
          checkpoint_path=checkpoint_path,
          master=self._config.evaluation_master,
          scaffold=estimator_spec.scaffold,
          eval_ops=update_op,
          final_ops=eval_dict,
          hooks=hooks,
          config=self._session_config)

      _write_dict_to_summary(
          output_dir=eval_dir,
          dictionary=eval_results,
          current_global_step=eval_results[ops.GraphKeys.GLOBAL_STEP])

    return eval_results

    def test_build_learning_rate_summaries(self):
        training.create_global_step()
        x = {'x': tf.placeholder(tf.float32, [2, 89])}
        y = tf.constant([[1], [1]])

        model = BaseModel(plx.Modes.TRAIN, graph_fn=self.get_dummy_graph_fn(),
                          loss_config=LossConfig(module='log_loss'),
                          optimizer_config=OptimizerConfig(module='adadelta',
                                                           decay_type='exponential_decay'),
                          model_type=BaseModel.Types.CLASSIFIER, eval_metrics_config=[],
                          summaries=['learning_rate'], name='test')

        model(x, y, None, None)

        # Only var are created
        summaries_names = list(get_tracked(collection=tf.GraphKeys.SUMMARIES_BY_NAMES).keys())
        assert len(summaries_names) == 1
        assert summaries_names[0] == 'learning_rate'

  def _create_global_step(self, graph):
    """Creates the global step tensor in graph.

    The global step tensor must be an integer type with name 'global_step' and
    be added to the collection ${tf.GraphKeys.GLOBAL_STEP}.

    Args:
      graph: The graph in which to create the global step tensor.

    Returns:
      The global step `Tensor`.
    """
    return training.create_global_step(graph)

 def test_do_not_stop_if_checkpoint_is_not_there(self):
   with ops.Graph().as_default():
     step = training.create_global_step()
     assign_ten = step.assign(10)
     no_op = control_flow_ops.no_op()
     hook = hooks_lib._StopAtCheckpointStepHook(
         model_dir=tempfile.mkdtemp(), last_step=10)
     with training.SingularMonitoredSession(hooks=[hook]) as mon_sess:
       mon_sess.raw_session().run(assign_ten)
       with test.mock.patch.object(time, 'sleep') as mock_sleep:
         mon_sess.run(no_op)
         self.assertTrue(mock_sleep.called)
       self.assertFalse(mon_sess.should_stop())

    def _evaluate_model(self, input_fn, hooks=None, checkpoint_path=None, name=''):
        # Check that model has been trained (if nothing has been set explicitly).
        if not checkpoint_path:
            latest_path = saver.latest_checkpoint(self._model_dir)
            if not latest_path:
                error_message = "Could not find trained model at {}.".format(self._model_dir)
                raise EstimatorNotTrainedError(error_message)
            checkpoint_path = latest_path

        # Setup output directory.
        eval_dir = os.path.join(self._model_dir, 'eval' if not name else 'eval_' + name)

        with ops.Graph().as_default() as g:
            random_seed.set_random_seed(self._config.tf_random_seed)
            global_step = training.create_global_step(g)
            features, labels = input_fn()

            estimator_spec = self._call_model_fn(features, labels, Modes.EVAL)
            if MetricKeys.LOSS in estimator_spec.eval_metric_ops:
                raise ValueError("Metric with name `{}` is not allowed, because Estimator "
                                 "already defines a default metric "
                                 "with the same name.".format(MetricKeys.LOSS))
            estimator_spec.eval_metric_ops[
                MetricKeys.LOSS] = metrics_lib.streaming_mean(estimator_spec.loss)
            update_op, eval_dict = self._extract_metric_update_ops(estimator_spec.eval_metric_ops)

            if ops.GraphKeys.GLOBAL_STEP in eval_dict:
                raise ValueError("Metric with name `global_step` is not allowed, because "
                                 "Estimator already defines a default metric with the same name.")
            eval_dict[ops.GraphKeys.GLOBAL_STEP] = global_step

            eval_results = evaluation._evaluate_once(
                checkpoint_path=checkpoint_path,
                master=self._config.evaluation_master,
                scaffold=estimator_spec.scaffold,
                eval_ops=update_op,
                final_ops=eval_dict,
                hooks=hooks,
                config=self._session_config)

            self._write_dict_to_summary(
                output_dir=eval_dir,
                dictionary=eval_results,
                current_global_step=eval_results[ops.GraphKeys.GLOBAL_STEP])

            return eval_results

  def export_savedmodel(
      self, export_dir_base, serving_input_receiver_fn,
      assets_extra=None,
      as_text=False,
      checkpoint_path=None):
    """Exports inference graph as a SavedModel into given dir.

    This method builds a new graph by first calling the
    serving_input_receiver_fn to obtain feature `Tensor`s, and then calling
    this `Estimator`'s model_fn to generate the model graph based on those
    features. It restores the given checkpoint (or, lacking that, the most
    recent checkpoint) into this graph in a fresh session.  Finally it creates
    a timestamped export directory below the given export_dir_base, and writes
    a `SavedModel` into it containing a single `MetaGraphDef` saved from this
    session.

    The exported `MetaGraphDef` will provide one `SignatureDef` for each
    element of the export_outputs dict returned from the model_fn, named using
    the same keys.  One of these keys is always
    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, indicating which
    signature will be served when a serving request does not specify one.
    For each signature, the outputs are provided by the corresponding
    `ExportOutput`s, and the inputs are always the input receivers provided by
    the serving_input_receiver_fn.

    Extra assets may be written into the SavedModel via the extra_assets
    argument.  This should be a dict, where each key gives a destination path
    (including the filename) relative to the assets.extra directory.  The
    corresponding value gives the full path of the source file to be copied.
    For example, the simple case of copying a single file without renaming it
    is specified as `{'my_asset_file.txt': '/path/to/my_asset_file.txt'}`.

    Args:
      export_dir_base: A string containing a directory in which to create
        timestamped subdirectories containing exported SavedModels.
      serving_input_receiver_fn: A function that takes no argument and
        returns a `ServingInputReceiver`.
      assets_extra: A dict specifying how to populate the assets.extra directory
        within the exported SavedModel, or `None` if no extra assets are needed.
      as_text: whether to write the SavedModel proto in text format.
      checkpoint_path: The checkpoint path to export.  If `None` (the default),
        the most recent checkpoint found within the model directory is chosen.

    Returns:
      The string path to the exported directory.

    Raises:
      ValueError: if no serving_input_receiver_fn is provided, no export_outputs
          are provided, or no checkpoint can be found.
    """
    if serving_input_receiver_fn is None:
      raise ValueError('serving_input_receiver_fn must be defined.')

    with ops.Graph().as_default() as g:
      training.create_global_step(g)
      random_seed.set_random_seed(self._config.tf_random_seed)
      serving_input_receiver = serving_input_receiver_fn()

      # Call the model_fn and collect the export_outputs.
      estimator_spec = self._call_model_fn(
          features=serving_input_receiver.features,
          labels=None,
          mode=model_fn_lib.ModeKeys.PREDICT)

      # Build the SignatureDefs from receivers and all outputs
      signature_def_map = build_all_signature_defs(
          serving_input_receiver.receiver_tensors,
          estimator_spec.export_outputs)

      if not checkpoint_path:
        # Locate the latest checkpoint
        checkpoint_path = saver.latest_checkpoint(self._model_dir)
      if not checkpoint_path:
        raise ValueError("Couldn't find trained model at %s." % self._model_dir)

      export_dir = get_timestamped_export_dir(export_dir_base)

      # TODO(soergel): Consider whether MonitoredSession makes sense here
      with tf_session.Session() as session:

        saver_for_restore = estimator_spec.scaffold.saver or saver.Saver(
            sharded=True)
        saver_for_restore.restore(session, checkpoint_path)

        # TODO(b/36111876): replace legacy_init_op with main_op mechanism
        # pylint: disable=protected-access
        local_init_op = (
            estimator_spec.scaffold.local_init_op or
            monitored_session.Scaffold._default_local_init_op())
        # pylint: enable=protected-access

        # Perform the export
        builder = saved_model_builder.SavedModelBuilder(export_dir)
        builder.add_meta_graph_and_variables(
            session, [tag_constants.SERVING],
            signature_def_map=signature_def_map,
            assets_collection=ops.get_collection(
                ops.GraphKeys.ASSET_FILEPATHS),
            legacy_init_op=local_init_op)
        builder.save(as_text)
#.........这里部分代码省略.........

  def _TestQuantize_Conv2dWithoutBatchNorm(self, activation, activation_op_name,
                                           with_bypass, delay):
    """Tests quantization: inputs -> Conv2d no batch norm -> Activation.

    Args:
      activation: Callable that returns an Operation, a factory method for the
        Activation.
      activation_op_name: String, name of the Activation operation.
      with_bypass: Bool, when true there is an extra connection added from
        inputs to just before Activation.
      delay: Int (optional), delay in number of steps until quantization starts.
    """
    graph = ops.Graph()
    with graph.as_default():
      training.create_global_step(graph)

      batch_size, height, width, depth = 5, 128, 128, 3
      inputs = array_ops.zeros((batch_size, height, width, depth))
      stride = 1 if with_bypass else 2
      out_depth = 3 if with_bypass else 32
      activation_fn = None if with_bypass else activation
      scope = 'test/test2' if with_bypass else 'test'
      node = conv2d(inputs, out_depth, [5, 5], stride=stride, padding='SAME',
                    weights_initializer=self._WeightInit(0.09),
                    activation_fn=activation_fn, scope=scope)
      if with_bypass:
        node = math_ops.add(inputs, node, name='test/Add')
        node = activation(node, name='test/' + activation_op_name)
      update_barrier = control_flow_ops.no_op(name='update_barrier')
      with ops.control_dependencies([update_barrier]):
        array_ops.identity(node, name='control_dependency')

      quantize.Quantize(graph, quant_delay=delay)
    quantization_node_name = 'FakeQuantWithMinMaxVars'
    weights_quant = graph.get_operation_by_name(scope + '/weights_quant/' +
                                                quantization_node_name)
    self.assertEqual(weights_quant.type, quantization_node_name)
    expected_inputs = [
        scope + '/weights_quant/Minimum', scope + '/weights_quant/Maximum',
        scope + '/weights/read'
    ]
    self._AssertInputOpsAre(weights_quant, expected_inputs)
    output_op_name = scope + '/convolution'
    self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name])

    if with_bypass:
      conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' +
                                               quantization_node_name)
      self.assertEqual(conv_quant.type, quantization_node_name)
      expected_inputs = [
          scope + '/conv_quant/min/read', scope + '/conv_quant/max/read',
          scope + '/BiasAdd'
      ]
      self._AssertInputOpsAre(conv_quant, expected_inputs)
      output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1'
                        if delay else 'test/Add')
      self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name])

    act_quant = graph.get_operation_by_name('test/act_quant/' +
                                            quantization_node_name)
    self.assertEqual(act_quant.type, quantization_node_name)

    expected_inputs = [
        'test/act_quant/min/read', 'test/act_quant/max/read',
        'test/' + activation_op_name
    ]
    self._AssertInputOpsAre(act_quant, expected_inputs)
    output_op_name = ('test/act_quant/delayed_quant/Switch_1'
                      if delay else 'control_dependency')
    self._AssertOutputGoesToOps(act_quant, graph, [output_op_name])

  def _testQuantize_DepthwiseConv2dWithBatchNorm(
      self, activation, activation_op_name, with_bypass, delay,
      fused_batch_norm, use_ema):
    """Tests quantization: inputs -> DWConv2d with batch norm -> Activation.

    Args:
      activation: Callable that returns an Operation, a factory method for the
        Activation.
      activation_op_name: String, name of the Activation operation.
      with_bypass: Bool, when true there is an extra connection added from
        inputs to just before Activation.
      delay: Int (optional), delay in number of steps until quantization starts.
      fused_batch_norm: Bool, when true use FusedBatchNorm.
      use_ema: Bool, when true uses EMA quantization for BN folded weights.
    """
    graph = ops.Graph()
    with graph.as_default():
      training.create_global_step(graph)

      batch_size, height, width, depth = 5, 128, 128, 3
      inputs = array_ops.zeros((batch_size, height, width, depth))
      stride = 1 if with_bypass else 2
      scope = 'test/test2' if with_bypass else 'test'
      node = separable_conv2d(
          inputs,
          None, [5, 5],
          stride=stride,
          depth_multiplier=1.0,
          padding='SAME',
          weights_initializer=self._WeightInit(0.09),
          activation_fn=None,
          normalizer_fn=batch_norm,
          normalizer_params=self._BatchNormParams(fused_batch_norm),
          scope=scope)

      # Manually add a bypass (optionaly) and an activation.
      if with_bypass:
        node = math_ops.add(inputs, node, name='test/Add')

      node = activation(node, name='test/' + activation_op_name)

      update_barrier = control_flow_ops.no_op(name='update_barrier')
      with ops.control_dependencies([update_barrier]):
        array_ops.identity(node, name='control_dependency')

      fold_batch_norms.FoldBatchNorms(graph)

      quantize.Quantize(
          graph, quant_delay=delay, quantize_folded_weights_use_ema=use_ema)
    quantization_node_name = 'FakeQuantWithMinMaxVars'
    weights_quant = graph.get_operation_by_name(scope + '/weights_quant/' +
                                                quantization_node_name)
    self.assertEqual(weights_quant.type, quantization_node_name)
    expected_inputs = [
        scope + '/weights_quant/' + ('min/read' if use_ema else 'Minimum'),
        scope + '/weights_quant/' + ('max/read' if use_ema else 'Maximum'),
        scope + '/mul_fold'
    ]
    self._AssertInputOpsAre(weights_quant, expected_inputs)
    output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1'
                              if delay and use_ema else '/depthwise_Fold')
    self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name])

    if with_bypass:
      conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' +
                                               quantization_node_name)
      self.assertEqual(conv_quant.type, quantization_node_name)
      expected_inputs = [
          scope + '/conv_quant/min/read', scope + '/conv_quant/max/read',
          scope + '/add_fold'
      ]
      self._AssertInputOpsAre(conv_quant, expected_inputs)
      output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1'
                        if delay else 'test/Add')
      self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name])

    act_quant = graph.get_operation_by_name('test/act_quant/' +
                                            quantization_node_name)
    self.assertEqual(act_quant.type, quantization_node_name)
    expected_inputs = [
        'test/act_quant/min/read', 'test/act_quant/max/read',
        'test/' + activation_op_name
    ]
    self._AssertInputOpsAre(act_quant, expected_inputs)
    output_op_name = ('test/act_quant/delayed_quant/Switch_1'
                      if delay else 'control_dependency')
    self._AssertOutputGoesToOps(act_quant, graph, [output_op_name])

  def _testQuantize_FCWithBatchNorm(self, activation, activation_op_name,
                                    with_bypass, delay, use_ema):
    """Tests quantization: inputs -> FC with batch norm -> Activation.

    Args:
      activation: Callable that returns an Operation, a factory method for the
        Activation.
      activation_op_name: String, name of the Activation operation.
      with_bypass: Bool, when true there is an extra connection added from
        inputs to just before Activation.
      delay: Int (optional), delay in number of steps until quantization starts.
      use_ema: Bool, when true uses EMA quantization for BN folded weights.
    """
    graph = ops.Graph()
    with graph.as_default():
      training.create_global_step(graph)

      batch_size, depth = 5, 256
      inputs = array_ops.zeros((batch_size, depth))
      out_depth = 256 if with_bypass else 128
      scope = 'test/test2' if with_bypass else 'test'
      node = fully_connected(inputs, out_depth,
                             weights_initializer=self._WeightInit(0.03),
                             activation_fn=None,
                             normalizer_fn=batch_norm,
                             normalizer_params=_DEFAULT_BATCH_NORM_PARAMS,
                             scope=scope)
      # Manually fold the batch norm.
      weights = graph.get_operation_by_name(scope + '/weights/read').outputs[0]
      bn_mult = (graph.get_operation_by_name(scope + '/BatchNorm/batchnorm/mul')
                 .outputs[0])
      mul_fold = math_ops.multiply(weights, bn_mult, name=scope + '/mul_fold')
      fc_fold = math_ops.matmul(inputs, mul_fold, name=scope + '/MatMul_Fold')
      bn_bias = (graph.get_operation_by_name(scope + '/BatchNorm/batchnorm/sub')
                 .outputs[0])
      add_fold = math_ops.add(fc_fold, bn_bias, name=scope + '/add_fold')
      # Manually add a bypass (optionaly) and an activation.
      if with_bypass:
        node = math_ops.add(inputs, add_fold, name='test/Add')
      else:
        node = add_fold
      node = activation(node, name='test/' + activation_op_name)

      update_barrier = control_flow_ops.no_op(name='update_barrier')
      with ops.control_dependencies([update_barrier]):
        array_ops.identity(node, name='control_dependency')

      quantize.Quantize(
          graph, quant_delay=delay, quantize_folded_weights_use_ema=use_ema)

    quantization_node_name = 'FakeQuantWithMinMaxVars'
    weights_quant = graph.get_operation_by_name(scope + '/weights_quant/' +
                                                quantization_node_name)
    self.assertEqual(weights_quant.type, quantization_node_name)
    expected_inputs = [
        scope + '/weights_quant/' + ('min/read' if use_ema else 'Minimum'),
        scope + '/weights_quant/' + ('max/read' if use_ema else 'Maximum'),
        scope + '/mul_fold'
    ]
    self._AssertInputOpsAre(weights_quant, expected_inputs)
    output_op_name = scope + ('/weights_quant/delayed_quant/Switch_1'
                              if delay and use_ema else '/MatMul_Fold')
    self._AssertOutputGoesToOps(weights_quant, graph, [output_op_name])

    if with_bypass:
      conv_quant = graph.get_operation_by_name(scope + '/conv_quant/' +
                                               quantization_node_name)
      self.assertEqual(conv_quant.type, quantization_node_name)
      expected_inputs = [
          scope + '/conv_quant/min/read', scope + '/conv_quant/max/read',
          scope + '/add_fold'
      ]
      self._AssertInputOpsAre(conv_quant, expected_inputs)
      output_op_name = (scope + '/conv_quant/delayed_quant/Switch_1'
                        if delay else 'test/Add')
      self._AssertOutputGoesToOps(conv_quant, graph, [output_op_name])

    act_quant = graph.get_operation_by_name('test/act_quant/' +
                                            quantization_node_name)
    self.assertEqual(act_quant.type, quantization_node_name)
    expected_inputs = [
        'test/act_quant/min/read', 'test/act_quant/max/read',
        'test/' + activation_op_name
    ]
    self._AssertInputOpsAre(act_quant, expected_inputs)
    output_op_name = ('test/act_quant/delayed_quant/Switch_1'
                      if delay else 'control_dependency')
    self._AssertOutputGoesToOps(act_quant, graph, [output_op_name])