(self, features, labels, mode, config)
return estimator_spec._replace(export_outputs=new_export_outputs)
Wraps `_call_model_fn` for `export_savedmodel`.
Wraps `_call_model_fn` for `export_savedmodel`.
Wraps _call_model_fn for export_savedmodel .
def _call_model_fn_for_inference(self, features, labels, mode, config):
    """Wraps `_call_model_fn` for `export_savedmodel`."""
    if mode != _REWRITE_FOR_INFERENCE_MODE:
      raise ValueError('mode must be {}; '
                       'got {}.'.format(_REWRITE_FOR_INFERENCE_MODE, mode))

    capture = _CapturedObject()

    def computation():
      """Compute tpu tensors used in export_outputs.

      Passed to rewrite_for_inference so that model_fn will be called under
      the rewriting contexts. Only tpu tensors are returned, but export_outputs
      and scaffold are captured.

      Returns:
         A list of Tensors used in export_outputs and not marked for
         outside_compilation.
      """
      # We should only call model fn once and it should be inside `computation`
      # so that building the graph will happen under `rewrite_for_inference`.
      mode = model_fn_lib.ModeKeys.PREDICT
      estimator_spec = self._call_model_fn(features, labels, mode, config)

      # We pick the TPU tensors out from `export_output` and later return them
      # from `computation` for rewriting.
      tensors_dict = collections.OrderedDict(
          (k, _export_output_to_tensors(v))
          for k, v in six.iteritems(estimator_spec.export_outputs))
      tensors = nest.flatten(tensors_dict)
      tpu_tensors = [t for t in tensors if t is not None]

      # We cannot return anything other than `tpu_tensors` here so we capture
      # the rest for later use.
      capture.capture((estimator_spec, tensors_dict, tensors))
      return tpu_tensors

    tpu_tensors_on_cpu = tpu.rewrite_for_inference(computation)
    estimator_spec, tensors_dict, tensors = capture.get()

    # Reconstruct `tensors`, but with `tpu_tensors` replaced with
    # `tpu_tensors_on_cpu`.
    new_tensors = []
    for t in tensors:
      if t is None:
        new_tensors.append(None)
      else:
        new_tensors.append(tpu_tensors_on_cpu.pop(0))

    # Reconstruct `tensors_dict`.
    new_tensors_dict = nest.pack_sequence_as(tensors_dict, new_tensors)
    # Reconstruct `export_outputs`.
    export_outputs = estimator_spec.export_outputs
    new_export_outputs = collections.OrderedDict(
        (k, _clone_export_output_with_tensors(export_outputs[k], v))
        for k, v in six.iteritems(new_tensors_dict))

    return estimator_spec._replace(export_outputs=new_export_outputs)
def _call_model_fn_for_inference ( self , features , labels , mode , config ) : if mode != _REWRITE_FOR_INFERENCE_MODE : raise ValueError ( 'mode must be {}; ' 'got {}.' . format ( _REWRITE_FOR_INFERENCE_MODE , mode ) ) capture = _CapturedObject ( ) def computation ( ) : """Compute tpu tensors used in export_outputs.

      Passed to rewrite_for_inference so that model_fn will be called under
      the rewriting contexts. Only tpu tensors are returned, but export_outputs
      and scaffold are captured.

      Returns:
         A list of Tensors used in export_outputs and not marked for
         outside_compilation.
      """ # We should only call model fn once and it should be inside `computation` # so that building the graph will happen under `rewrite_for_inference`. mode = model_fn_lib . ModeKeys . PREDICT estimator_spec = self . _call_model_fn ( features , labels , mode , config ) # We pick the TPU tensors out from `export_output` and later return them # from `computation` for rewriting. tensors_dict = collections . OrderedDict ( ( k , _export_output_to_tensors ( v ) ) for k , v in six . iteritems ( estimator_spec . export_outputs ) ) tensors = nest . flatten ( tensors_dict ) tpu_tensors = [ t for t in tensors if t is not None ] # We cannot return anything other than `tpu_tensors` here so we capture # the rest for later use. capture . capture ( ( estimator_spec , tensors_dict , tensors ) ) return tpu_tensors tpu_tensors_on_cpu = tpu . rewrite_for_inference ( computation ) estimator_spec , tensors_dict , tensors = capture . get ( ) # Reconstruct `tensors`, but with `tpu_tensors` replaced with # `tpu_tensors_on_cpu`. new_tensors = [ ] for t in tensors : if t is None : new_tensors . append ( None ) else : new_tensors . append ( tpu_tensors_on_cpu . pop ( 0 ) ) # Reconstruct `tensors_dict`. new_tensors_dict = nest . pack_sequence_as ( tensors_dict , new_tensors ) # Reconstruct `export_outputs`. export_outputs = estimator_spec . export_outputs new_export_outputs = collections . OrderedDict ( ( k , _clone_export_output_with_tensors ( export_outputs [ k ] , v ) ) for k , v in six . iteritems ( new_tensors_dict ) ) return estimator_spec . _replace ( export_outputs = new_export_outputs )