def test_export_strategies_reset(self):
    est = TestEstimator()
    export_strategy_1 = saved_model_export_utils.make_export_strategy(
        est, 'export_input_1', exports_to_keep=None)

    ex = experiment.Experiment(
        est,
        train_input_fn='train_input',
        eval_input_fn='eval_input',
        eval_metrics='eval_metrics',
        train_steps=100,
        eval_steps=100,
        export_strategies=[export_strategy_1])
    ex.train_and_evaluate()
    self.assertEqual(1, est.export_count)

    # After reset with empty list (None), the count does not change and the user
    # provided export strategy list should remain intact.
    old_es = ex.reset_export_strategies()
    ex.train_and_evaluate()
    self.assertAllEqual([export_strategy_1], old_es)
    self.assertEqual(1, est.export_count)

    # After reset with list, the count should increase with the number of items.
    export_strategy_2 = saved_model_export_utils.make_export_strategy(
        est, 'export_input_2', exports_to_keep=None)
    export_strategy_3 = saved_model_export_utils.make_export_strategy(
        est, 'export_input_3', exports_to_keep=None)

    old_es = ex.reset_export_strategies([export_strategy_2, export_strategy_3])
    ex.train_and_evaluate()
    self.assertAllEqual([], old_es)
    self.assertEqual(3, est.export_count)

def _make_experiment_fn(output_dir):
  """Creates experiment for DNNBoostedTreeCombinedRegressor."""
  (x_train, y_train), (x_test,
                       y_test) = tf.keras.datasets.boston_housing.load_data()

  train_input_fn = tf.compat.v1.estimator.inputs.numpy_input_fn(
      x={"x": x_train},
      y=y_train,
      batch_size=FLAGS.batch_size,
      num_epochs=None,
      shuffle=True)
  eval_input_fn = tf.compat.v1.estimator.inputs.numpy_input_fn(
      x={"x": x_test}, y=y_test, num_epochs=1, shuffle=False)

  feature_columns = [
      feature_column.real_valued_column("x", dimension=_BOSTON_NUM_FEATURES)
  ]
  feature_spec = tf.contrib.layers.create_feature_spec_for_parsing(
      feature_columns)
  serving_input_fn = input_fn_utils.build_parsing_serving_input_fn(feature_spec)
  export_strategies = [
      saved_model_export_utils.make_export_strategy(serving_input_fn)]
  return tf.contrib.learn.Experiment(
      estimator=_get_estimator(output_dir, feature_columns),
      train_input_fn=train_input_fn,
      eval_input_fn=eval_input_fn,
      train_steps=None,
      eval_steps=FLAGS.num_eval_steps,
      eval_metrics=None,
      export_strategies=export_strategies)

def make_custom_export_strategy(name,
                                convert_fn,
                                feature_columns,
                                export_input_fn):
  """Makes custom exporter of GTFlow tree format.

  Args:
    name: A string, for the name of the export strategy.
    convert_fn: A function that converts the tree proto to desired format and
      saves it to the desired location. Can be None to skip conversion.
    feature_columns: A list of feature columns.
    export_input_fn: A function that takes no arguments and returns an
      `InputFnOps`.

  Returns:
    An `ExportStrategy`.
  """
  base_strategy = saved_model_export_utils.make_export_strategy(
      serving_input_fn=export_input_fn)
  input_fn = export_input_fn()
  (sorted_feature_names, dense_floats, sparse_float_indices, _, _,
   sparse_int_indices, _, _) = gbdt_batch.extract_features(
       input_fn.features, feature_columns)

  def export_fn(estimator, export_dir, checkpoint_path=None, eval_result=None):
    """A wrapper to export to SavedModel, and convert it to other formats."""
    result_dir = base_strategy.export(estimator, export_dir,
                                      checkpoint_path,
                                      eval_result)
    with ops.Graph().as_default() as graph:
      with tf_session.Session(graph=graph) as sess:
        saved_model_loader.load(
            sess, [tag_constants.SERVING], result_dir)
        # Note: This is GTFlow internal API and might change.
        ensemble_model = graph.get_operation_by_name(
            "ensemble_model/TreeEnsembleSerialize")
        _, dfec_str = sess.run(ensemble_model.outputs)
        dtec = tree_config_pb2.DecisionTreeEnsembleConfig()
        dtec.ParseFromString(dfec_str)
        # Export the result in the same folder as the saved model.
        if convert_fn:
          convert_fn(dtec, sorted_feature_names,
                     len(dense_floats),
                     len(sparse_float_indices),
                     len(sparse_int_indices), result_dir, eval_result)
        feature_importances = _get_feature_importances(
            dtec, sorted_feature_names,
            len(dense_floats),
            len(sparse_float_indices), len(sparse_int_indices))
        sorted_by_importance = sorted(
            feature_importances.items(), key=lambda x: -x[1])
        assets_dir = os.path.join(result_dir, "assets.extra")
        gfile.MakeDirs(assets_dir)
        with gfile.GFile(os.path.join(assets_dir, "feature_importances"),
                         "w") as f:
          f.write("\n".join("%s, %f" % (k, v) for k, v in sorted_by_importance))
    return result_dir

  return export_strategy.ExportStrategy(
      name, export_fn, strip_default_attrs=True)

 def test_train_and_evaluate(self):
   for est in self._estimators_for_tests():
     eval_metrics = 'eval_metrics' if not isinstance(
         est, core_estimator.Estimator) else None
     noop_hook = _NoopHook()
     export_strategy = saved_model_export_utils.make_export_strategy(
         est,
         None if isinstance(est, core_estimator.Estimator) else 'export_input',
         exports_to_keep=None)
     ex = experiment.Experiment(
         est,
         train_input_fn='train_input',
         eval_input_fn='eval_input',
         eval_metrics=eval_metrics,
         eval_hooks=[noop_hook],
         train_steps=100,
         eval_steps=100,
         export_strategies=export_strategy)
     ex.train_and_evaluate()
     self.assertEqual(1, est.fit_count)
     self.assertEqual(1, est.eval_count)
     self.assertEqual(1, est.export_count)
     self.assertEqual(1, len(est.monitors))
     self.assertEqual([noop_hook], est.eval_hooks)
     self.assertTrue(isinstance(est.monitors[0],
                                session_run_hook.SessionRunHook))

  def test_continuous_train_and_eval_with_predicate_fn(self):
    for est in self._estimators_for_tests(eval_dict={'global_step': 100}):
      eval_metrics = 'eval_metrics' if not isinstance(
          est, core_estimator.Estimator) else None
      export_strategy = saved_model_export_utils.make_export_strategy(
          est,
          None if isinstance(est, core_estimator.Estimator) else 'export_input',
          exports_to_keep=None)
      ex = experiment.Experiment(
          est,
          train_input_fn='train_input',
          eval_input_fn='eval_input',
          eval_metrics=eval_metrics,
          train_steps=100000000000,  # a value will make `ex` never stops.
          eval_steps=100,
          export_strategies=export_strategy)

      def predicate_fn(eval_result):
        del eval_result  # unused. for fn signature.
        return False

      ex.continuous_train_and_eval(continuous_eval_predicate_fn=predicate_fn)
      self.assertEqual(0, est.fit_count)
      self.assertEqual(0, est.eval_count)
      self.assertEqual(1, est.export_count)

def main(unused_argv):

  # Load training and eval data
  mnist = read_data_sets(FLAGS.data_dir,
      source_url=FLAGS.datasource_url)

  train_data = mnist.train.images  # Returns np.array
  train_labels = np.asarray(mnist.train.labels, dtype=np.int32)
  eval_data = mnist.test.images  # Returns np.array
  eval_labels = np.asarray(mnist.test.labels, dtype=np.int32)


  def serving_input_receiver_fn():
      feature_tensor = tf.placeholder(tf.float32, [None, 784])
      return tf.estimator.export.ServingInputReceiver({'x': feature_tensor}, {'x': feature_tensor})

  learn_runner.run(
      generate_experiment_fn(
          min_eval_frequency=1,
          train_steps=FLAGS.num_steps,
          eval_steps=FLAGS.eval_steps,
          export_strategies=[saved_model_export_utils.make_export_strategy(
              serving_input_receiver_fn,
              exports_to_keep=1
          )]
      ),
      run_config = tf.contrib.learn.RunConfig().replace(model_dir=FLAGS.job_dir, save_checkpoints_steps=1000),
      hparams=hparam.HParams(dataset=mnist.train, eval_data=eval_data, eval_labels=eval_labels),
  )

  def test_checkpoint_and_export(self):
    model_dir = tempfile.mkdtemp()
    config = run_config_lib.RunConfig(save_checkpoints_steps=3)
    est = dnn.DNNClassifier(
        n_classes=3,
        feature_columns=[
            feature_column.real_valued_column('feature', dimension=4)
        ],
        hidden_units=[3, 3],
        model_dir=model_dir,
        config=config)

    exp_strategy = saved_model_export_utils.make_export_strategy(
        est, 'export_input', exports_to_keep=None)

    ex = experiment.Experiment(
        est,
        train_input_fn=test_data.iris_input_multiclass_fn,
        eval_input_fn=test_data.iris_input_multiclass_fn,
        export_strategies=(exp_strategy,),
        train_steps=8,
        checkpoint_and_export=True,
        eval_delay_secs=0)

    with test.mock.patch.object(ex, '_maybe_export'):
      with test.mock.patch.object(ex, '_call_evaluate'):
        ex.train_and_evaluate()
        # Eval and export are called after steps 1, 4, 7, and 8 (after training
        # is completed).
        self.assertEqual(ex._maybe_export.call_count, 4)
        self.assertEqual(ex._call_evaluate.call_count, 4)

 def test_continuous_train_and_eval(self):
   for est in self._estimators_for_tests(eval_dict={'global_step': 100}):
     if isinstance(est, core_estimator.Estimator):
       eval_metrics = None
       saving_listeners = 'saving_listeners'
     else:
       eval_metrics = 'eval_metrics'
       saving_listeners = None
     noop_hook = _NoopHook()
     export_strategy = saved_model_export_utils.make_export_strategy(
         est,
         None if isinstance(est, core_estimator.Estimator) else 'export_input',
         exports_to_keep=None)
     ex = experiment.Experiment(
         est,
         train_input_fn='train_input',
         eval_input_fn='eval_input',
         eval_metrics=eval_metrics,
         eval_hooks=[noop_hook],
         train_steps=100,
         eval_steps=100,
         export_strategies=export_strategy,
         saving_listeners=saving_listeners)
     ex.continuous_train_and_eval()
     self.assertEqual(1, est.fit_count)
     self.assertEqual(1, est.eval_count)
     self.assertEqual(1, est.export_count)
     self.assertEqual([noop_hook], est.eval_hooks)

 def _export_strategy():
     if self.saves_training():
         return [saved_model_export_utils.make_export_strategy(
             serving_input_fn=_serving_input_fn,
             default_output_alternative_key=None,
             exports_to_keep=1)]
     logger.warn("serving_input_fn not specified, model NOT saved, use checkpoints to reconstruct")
     return None

  def test_export_strategies_reset(self):
    for est in self._estimators_for_tests():
      eval_metrics = 'eval_metrics' if not isinstance(
          est, core_estimator.Estimator) else None
      export_strategy_1 = saved_model_export_utils.make_export_strategy(
          est,
          None if isinstance(est, core_estimator.Estimator) else 'export_1',
          exports_to_keep=None)

      ex = experiment.Experiment(
          est,
          train_input_fn='train_input',
          eval_input_fn='eval_input',
          eval_metrics=eval_metrics,
          train_steps=100,
          eval_steps=100,
          export_strategies=(export_strategy_1,))
      ex.train_and_evaluate()
      self.assertEqual(1, est.export_count)

      # After reset with empty list (None), the count does not change and the
      # user provided export strategy list should remain intact.
      old_es = ex.reset_export_strategies()
      ex.train_and_evaluate()
      self.assertAllEqual([export_strategy_1], old_es)
      self.assertEqual(1, est.export_count)

      # After reset with list, the count should increase with the number of
      # items.
      export_strategy_2 = saved_model_export_utils.make_export_strategy(
          est,
          None if isinstance(est, core_estimator.Estimator) else 'export_2',
          exports_to_keep=None)
      export_strategy_3 = saved_model_export_utils.make_export_strategy(
          est,
          None if isinstance(est, core_estimator.Estimator) else 'export_3',
          exports_to_keep=None)

      old_es = ex.reset_export_strategies(
          [export_strategy_2, export_strategy_3])
      ex.train_and_evaluate()
      self.assertAllEqual([], old_es)
      self.assertEqual(3, est.export_count)

 def test_make_export_strategy(self):
   """Only tests that an ExportStrategy instance is created."""
   def _serving_input_fn():
     return array_ops.constant([1]), None
   export_strategy = saved_model_export_utils.make_export_strategy(
       serving_input_fn=_serving_input_fn,
       default_output_alternative_key="default",
       assets_extra={"from/path": "to/path"},
       as_text=False,
       exports_to_keep=5)
   self.assertTrue(
       isinstance(export_strategy, export_strategy_lib.ExportStrategy))

 def test_test(self):
   for est in self._estimators_for_tests():
     exp_strategy = saved_model_export_utils.make_export_strategy(
         est, 'export_input', exports_to_keep=None)
     ex = experiment.Experiment(
         est,
         train_input_fn='train_input',
         eval_input_fn='eval_input',
         export_strategies=(exp_strategy,))
     ex.test()
     self.assertEqual(1, est.fit_count)
     self.assertEqual(1, est.eval_count)
     self.assertEqual(1, est.export_count)

 def test_test(self):
   est = TestEstimator()
   exp_strategy = saved_model_export_utils.make_export_strategy(
       est, 'export_input', exports_to_keep=None)
   ex = experiment.Experiment(
       est,
       train_input_fn='train_input',
       eval_input_fn='eval_input',
       export_strategies=[exp_strategy])
   ex.test()
   self.assertEqual(1, est.fit_count)
   self.assertEqual(1, est.eval_count)
   self.assertEqual(1, est.export_count)

 def test_default_output_alternative_key_core_estimator(self):
   est = TestCoreEstimator()
   export_strategy = saved_model_export_utils.make_export_strategy(
       est,
       default_output_alternative_key='export_key',
       exports_to_keep=None)
   ex = experiment.Experiment(
       est,
       train_input_fn='train_input',
       eval_input_fn='eval_input',
       train_steps=100,
       eval_steps=100,
       export_strategies=export_strategy)
   with self.assertRaisesRegexp(
       ValueError, 'default_output_alternative_key is not supported'):
     ex.train_and_evaluate()

  def _experiment_fn(output_dir):
    return tflearn.Experiment(
        get_model(output_dir, nbuckets, hidden_units, learning_rate),
        train_input_fn=read_dataset(traindata, mode=tf.contrib.learn.ModeKeys.TRAIN, num_training_epochs=num_training_epochs, batch_size=batch_size),
        eval_input_fn=read_dataset(evaldata),
        export_strategies=[saved_model_export_utils.make_export_strategy(
            serving_input_fn,
            default_output_alternative_key=None,
            exports_to_keep=1
        )],
        eval_metrics = {
	    'rmse' : tflearn.MetricSpec(metric_fn=my_rmse, prediction_key='probabilities'),
            'training/hptuning/metric' : tflearn.MetricSpec(metric_fn=my_rmse, prediction_key='probabilities')
        },
        min_eval_frequency = 100,
        **args
    )

 def test_train_and_evaluate(self):
   est = TestEstimator()
   export_strategy = saved_model_export_utils.make_export_strategy(
       est, 'export_input')
   ex = experiment.Experiment(
       est,
       train_input_fn='train_input',
       eval_input_fn='eval_input',
       eval_metrics='eval_metrics',
       train_steps=100,
       eval_steps=100,
       export_strategies=export_strategy)
   ex.train_and_evaluate()
   self.assertEquals(1, est.fit_count)
   self.assertEquals(1, est.eval_count)
   self.assertEquals(1, est.export_count)
   self.assertEquals(1, len(est.monitors))
   self.assertTrue(isinstance(est.monitors[0], monitors.ValidationMonitor))

  def _experiment_fn(output_dir):
    train_input_fn = generate_input_fn(train_file)
    eval_input_fn = generate_input_fn(test_file)
    my_model = build_estimator(model_type=model_type, 
                               model_dir=output_dir)

    experiment = tf.contrib.learn.Experiment(
      my_model,
      train_input_fn=train_input_fn,
      eval_input_fn=eval_input_fn,
      train_steps=1000
      ,
      export_strategies=[saved_model_export_utils.make_export_strategy(
        serving_input_fn,
        default_output_alternative_key=None
      )]
    )
    return experiment

def experiment_fn(output_dir):
    # run experiment
    return tflearn.Experiment(
        tflearn.Estimator(model_fn=cnn_model, model_dir=output_dir),
        train_input_fn=get_train(),
        eval_input_fn=get_valid(),
        eval_metrics={
            'acc': tflearn.MetricSpec(
                metric_fn=metrics.streaming_accuracy, prediction_key='class'
            )
        },
        export_strategies=[saved_model_export_utils.make_export_strategy(
            serving_input_fn,
            default_output_alternative_key=None,
            exports_to_keep=1
        )],
        train_steps = TRAIN_STEPS
    )

def make_custom_export_strategy(name, convert_fn, feature_columns,
                                export_input_fn):
  """Makes custom exporter of GTFlow tree format.

  Args:
    name: A string, for the name of the export strategy.
    convert_fn: A function that converts the tree proto to desired format and
      saves it to the desired location.
    feature_columns: A list of feature columns.
    export_input_fn: A function that takes no arguments and returns an
      `InputFnOps`.

  Returns:
    An `ExportStrategy`.
  """
  base_strategy = saved_model_export_utils.make_export_strategy(
      serving_input_fn=export_input_fn)
  input_fn = export_input_fn()
  (sorted_feature_names, dense_floats, sparse_float_indices, _, _,
   sparse_int_indices, _, _) = gbdt_batch.extract_features(
       input_fn.features, feature_columns)

  def export_fn(estimator, export_dir, checkpoint_path=None, eval_result=None):
    """A wrapper to export to SavedModel, and convert it to other formats."""
    result_dir = base_strategy.export(estimator, export_dir,
                                      checkpoint_path,
                                      eval_result)
    with ops.Graph().as_default() as graph:
      with tf_session.Session(graph=graph) as sess:
        saved_model_loader.load(
            sess, [tag_constants.SERVING], result_dir)
        # Note: This is GTFlow internal API and might change.
        ensemble_model = graph.get_operation_by_name(
            "ensemble_model/TreeEnsembleSerialize")
        _, dfec_str = sess.run(ensemble_model.outputs)
        dtec = tree_config_pb2.DecisionTreeEnsembleConfig()
        dtec.ParseFromString(dfec_str)
        # Export the result in the same folder as the saved model.
        convert_fn(dtec, sorted_feature_names, len(dense_floats),
                   len(sparse_float_indices), len(sparse_int_indices),
                   result_dir, eval_result)
    return result_dir
  return export_strategy.ExportStrategy(name, export_fn)

 def experiment_fn(output_dir):
    get_train = model.read_dataset(train_data_paths, mode=tf.contrib.learn.ModeKeys.TRAIN)
    get_valid = model.read_dataset(eval_data_paths, mode=tf.contrib.learn.ModeKeys.EVAL)
    # run experiment
    return tflearn.Experiment(
        tflearn.Estimator(model_fn=model.simple_rnn, model_dir=output_dir),
        train_input_fn=get_train,
        eval_input_fn=get_valid,
        eval_metrics={
            'rmse': tflearn.MetricSpec(
                metric_fn=tf.contrib.metrics.streaming_root_mean_squared_error
            )
        },
        export_strategies=[saved_model_export_utils.make_export_strategy(
            model.serving_input_fn,
            default_output_alternative_key=None,
            exports_to_keep=1
        )],
        **experiment_args
    )