text: # Copyright 2016 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.
  # ==============================================================================
  """A tf.learn implementation of tensor_forest (extremely random forests)."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from tensorflow.contrib import framework as contrib_framework
  from tensorflow.contrib.framework import deprecated_arg_values
  from tensorflow.contrib.framework.python.framework import experimental
  from tensorflow.contrib.learn.python.learn import evaluable
  from tensorflow.contrib.learn.python.learn import trainable
  
  from tensorflow.contrib.learn.python.learn.estimators import estimator
  from tensorflow.contrib.learn.python.learn.utils import export
  
  from tensorflow.contrib.tensor_forest.client import eval_metrics
  from tensorflow.contrib.tensor_forest.data import data_ops
  from tensorflow.contrib.tensor_forest.python import tensor_forest
  
  from tensorflow.python.framework import dtypes
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import state_ops
  from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.training import session_run_hook
  
  
  KEYS_NAME = 'keys'
  LOSS_NAME = 'rf_training_loss'
  
  
  def _assert_float32(tensors):
    """Assert all tensors are float32.
  
    Args:
      tensors: `Tensor` or `dict` of `Tensor` objects.
  
    Raises:
      TypeError: if any tensor is not float32.
    """
    if not isinstance(tensors, dict):
      tensors = [tensors]
    else:
      tensors = tensors.values()
    for tensor in tensors:
      if tensor.dtype.base_dtype != dtypes.float32:
        raise TypeError('Expected dtype=float32, %s.' % tensor)
  
  
  class TensorForestLossHook(session_run_hook.SessionRunHook):
    """Monitor to request stop when loss stops decreasing."""
  
    def __init__(self, early_stopping_rounds):
      self.early_stopping_rounds = early_stopping_rounds
      self.min_loss = None
      self.last_step = -1
      # self.steps records the number of steps for which the loss has been
      # non-decreasing
      self.steps = 0
  
    def before_run(self, run_context):
      return session_run_hook.SessionRunArgs(
          {'global_step': contrib_framework.get_global_step(),
           'current_loss': run_context.session.graph.get_operation_by_name(
               LOSS_NAME).outputs[0]})
  
    def after_run(self, run_context, run_values):
      current_loss = run_values.results['current_loss']
      current_step = run_values.results['global_step']
      self.steps += 1
      # Gaurd against the global step going backwards, which might happen
      # if we recover from something.
      if self.last_step == -1 or self.last_step > current_step:
        logging.info('TensorForestLossHook resetting last_step.')
        self.last_step = current_step
        self.steps = 0
        self.min_loss = None
        return
  
      self.last_step = current_step
      if self.min_loss is None or current_loss < self.min_loss:
        self.min_loss = current_loss
        self.steps = 0
      if self.steps > self.early_stopping_rounds:
        logging.info('TensorForestLossHook requesting stop.')
        run_context.request_stop()
  
  
  def get_model_fn(params, graph_builder_class, device_assigner,
                   weights_name=None, keys_name=None):
    """Return a model function given a way to construct a graph builder."""
    def _model_fn(features, labels):
      """Function that returns predictions, training loss, and training op."""
      weights = None
      keys = None
      if weights_name and weights_name in features:
        weights = features.pop(weights_name)
      if keys_name and keys_name in features:
        keys = features.pop(keys_name)
      processed_features, spec = data_ops.ParseDataTensorOrDict(features)
      _assert_float32(processed_features)
      if labels is not None:
        labels = data_ops.ParseLabelTensorOrDict(labels)
        _assert_float32(labels)
  
      graph_builder = graph_builder_class(params, device_assigner=device_assigner)
      inference = {eval_metrics.INFERENCE_PROB_NAME:
                   graph_builder.inference_graph(processed_features,
                                                 data_spec=spec)}
      if not params.regression:
        inference[eval_metrics.INFERENCE_PRED_NAME] = math_ops.argmax(
            inference[eval_metrics.INFERENCE_PROB_NAME], 1)
      if keys:
        inference[KEYS_NAME] = keys
  
      # labels might be None if we're doing prediction (which brings up the
      # question of why we force everything to adhere to a single model_fn).
      training_loss = None
      training_graph = None
      if labels is not None:
        training_loss = graph_builder.training_loss(processed_features, labels,
                                                    data_spec=spec,
                                                    name=LOSS_NAME)
        training_graph = control_flow_ops.group(
            graph_builder.training_graph(
                processed_features, labels, data_spec=spec,
                input_weights=weights),
            state_ops.assign_add(contrib_framework.get_global_step(), 1))
      # Put weights back in
      if weights is not None:
        features[weights_name] = weights
      return (inference, training_loss, training_graph)
    return _model_fn
  
  
  class TensorForestEstimator(evaluable.Evaluable, trainable.Trainable):
    """An estimator that can train and evaluate a random forest.
  
    Example:
  
    ```python
    params = tf.contrib.tensor_forest.python.tensor_forest.ForestHParams(
        num_classes=2, num_features=40, num_trees=10, max_nodes=1000)
  
    # Estimator using the default graph builder.
    estimator = TensorForestEstimator(params, model_dir=model_dir)
  
    # Or estimator using TrainingLossForest as the graph builder.
    estimator = TensorForestEstimator(
        params, graph_builder_class=tensor_forest.TrainingLossForest,
        model_dir=model_dir)
  
    # Input builders
    def input_fn_train: # returns x, y
      ...
    def input_fn_eval: # returns x, y
      ...
    estimator.fit(input_fn=input_fn_train)
    estimator.evaluate(input_fn=input_fn_eval)
    estimator.predict(x=x)
    ```
    """
  
    def __init__(self, params, device_assigner=None, model_dir=None,
                 graph_builder_class=tensor_forest.RandomForestGraphs,
                 config=None, weights_name=None, keys_name=None,
                 feature_engineering_fn=None, early_stopping_rounds=100):
  
      """Initializes a TensorForestEstimator instance.
  
      Args:
        params: ForestHParams object that holds random forest hyperparameters.
          These parameters will be passed into `model_fn`.
        device_assigner: An `object` instance that controls how trees get
          assigned to devices. If `None`, will use
          `tensor_forest.RandomForestDeviceAssigner`.
        model_dir: Directory to save model parameters, graph, etc. To continue
          training a previously saved model, load checkpoints saved to this
          directory into an estimator.
        graph_builder_class: An `object` instance that defines how TF graphs for
          random forest training and inference are built. By default will use
          `tensor_forest.RandomForestGraphs`.
        config: `RunConfig` object to configure the runtime settings.
        weights_name: A string defining feature column name representing
          weights. Will be multiplied by the loss of the example. Used to
          downweight or boost examples during training.
        keys_name: A string defining feature column name representing example
          keys. Used by `predict_with_keys` method.
        feature_engineering_fn: Feature engineering function. Takes features and
          labels which are the output of `input_fn` and returns features and
          labels which will be fed into the model.
        early_stopping_rounds: Allows training to terminate early if the forest is
          no longer growing. 100 by default.
  
      Returns:
        A `TensorForestEstimator` instance.
      """
      self.params = params.fill()
      self.graph_builder_class = graph_builder_class
      self.early_stopping_rounds = early_stopping_rounds
      self.weights_name = weights_name
      self._estimator = estimator.Estimator(
          model_fn=get_model_fn(params, graph_builder_class, device_assigner,
                                weights_name=weights_name, keys_name=keys_name),
          model_dir=model_dir,
          config=config,
          feature_engineering_fn=feature_engineering_fn)
      self._skcompat = estimator.SKCompat(self._estimator)
  
    @property
    def model_dir(self):
      """See evaluable.Evaluable."""
      return self._estimator.model_dir
  
    def evaluate(
        self, x=None, y=None, input_fn=None, batch_size=None,
        steps=None, metrics=None, name=None, checkpoint_path=None):
      """See evaluable.Evaluable."""
      if x is not None and y is not None:
        return self._skcompat.score(x, y, batch_size=batch_size, steps=steps,
                                    metrics=metrics)
      elif input_fn is not None:
        return self._estimator.evaluate(input_fn=input_fn, steps=steps,
                                        metrics=metrics, name=name,
                                        checkpoint_path=checkpoint_path)
      else:
        raise ValueError(
            'evaluate: Must provide either both x and y or input_fn.')
  
    def fit(self, x=None, y=None, input_fn=None, steps=None, batch_size=None,
            monitors=None, max_steps=None):
      """See trainable.Trainable."""
      if not monitors:
        monitors = [TensorForestLossHook(self.early_stopping_rounds)]
      if x is not None and y is not None:
        self._skcompat.fit(x, y, batch_size=batch_size, steps=steps,
                           max_steps=max_steps, monitors=monitors)
      elif input is not None:
        self._estimator.fit(input_fn=input_fn, steps=steps, monitors=monitors,
                            max_steps=max_steps)
      else:
        raise ValueError('fit: Must provide either both x and y or input_fn.')
  
    def predict_proba(
        self, x=None, input_fn=None, batch_size=None):
      """Returns prediction probabilities for given features (classification).
  
      Args:
        x: features.
        input_fn: Input function. If set, x and y must be None.
        batch_size: Override default batch size.
  
      Returns:
        Numpy array of predicted probabilities (or an iterable of predicted
        probabilities if as_iterable is True).
  
      Raises:
        ValueError: If both or neither of x and input_fn were given.
      """
      if x is not None:
        results = self._skcompat.predict(x, batch_size=batch_size)
        return results[eval_metrics.INFERENCE_PROB_NAME]
      else:
        results = self._estimator.predict(input_fn=input_fn, as_iterable=True)
        return (x[eval_metrics.INFERENCE_PROB_NAME] for x in results)
  
    def predict(
        self, x=None, input_fn=None, axis=None, batch_size=None):
      """Returns predictions for given features.
  
      Args:
        x: features.
        input_fn: Input function. If set, x must be None.
        axis: Axis on which to argmax (for classification).
              Last axis is used by default.
        batch_size: Override default batch size.
  
      Returns:
        Numpy array of predicted classes or regression values (or an iterable of
        predictions if as_iterable is True).
      """
      predict_name = (eval_metrics.INFERENCE_PROB_NAME if self.params.regression
                      else eval_metrics.INFERENCE_PRED_NAME)
      if x is not None:
        results = self._skcompat.predict(x, batch_size=batch_size)
        return results[predict_name]
      else:
        results = self._estimator.predict(input_fn=input_fn, as_iterable=True)
        return (x[predict_name] for x in results)
  
    def predict_with_keys(
        self, x=None, input_fn=None, axis=None, batch_size=None):
      """Same as predict but also returns the example keys."""
      predict_name = (eval_metrics.INFERENCE_PROB_NAME if self.params.regression
                      else eval_metrics.INFERENCE_PRED_NAME)
      if x is not None:
        results = self._skcompat.predict(x, batch_size=batch_size)
        return results[predict_name]
      else:
        results = self._estimator.predict(input_fn=input_fn, as_iterable=True)
        return ((x[predict_name], x.get(KEYS_NAME, None)) for x in results)
  
    def export(self,
               export_dir,
               input_fn,
               signature_fn=None,
               input_feature_key=None,
               default_batch_size=1):
      """See BaseEstimator.export."""
      # Reset model function with basic device assigner.
      # Servo doesn't support distributed inference
      # but it will try to respect device assignments if they're there.
      # pylint: disable=protected-access
      orig_model_fn = self._estimator._model_fn
      self._estimator._model_fn = get_model_fn(
          self.params, self.graph_builder_class,
          tensor_forest.RandomForestDeviceAssigner(),
          weights_name=self.weights_name)
      result = self._estimator.export(
          export_dir=export_dir,
          input_fn=input_fn,
          input_feature_key=input_feature_key,
          use_deprecated_input_fn=False,
          signature_fn=(signature_fn or
                        (export.regression_signature_fn
                         if self.params.regression else
                         export.classification_signature_fn_with_prob)),
          default_batch_size=default_batch_size,
          prediction_key=eval_metrics.INFERENCE_PROB_NAME)
      self._estimator._model_fn = orig_model_fn
      # pylint: enable=protected-access
      return result
  
    @experimental
    def export_savedmodel(self,
                          export_dir_base,
                          input_fn,
                          default_output_alternative_key=None,
                          assets_extra=None,
                          as_text=False,
                          exports_to_keep=None):
      return self._estimator.export_savedmodel(
          export_dir_base,
          input_fn,
          default_output_alternative_key=default_output_alternative_key,
          assets_extra=assets_extra,
          as_text=as_text,
          exports_to_keep=exports_to_keep)
  