How to use the polyaxon.estimators.Estimator function in polyaxon

def test_default_model_dir(self):
        with test.mock.patch.object(tempfile, 'mkdtemp', return_value=_TMP_DIR):
            est = Estimator(model_fn=self.get_dummy_model_fn())
            self.assertIn(_TMP_DIR, est.config.model_dir)
            self.assertIn(_TMP_DIR, est.model_dir)

def test_min_eval_frequency_defaults(self):
        def dummy_model_fn(features, labels):  # pylint: disable=unused-argument
            pass

        # The default value when model_dir is on GCS is 1000
        estimator = Estimator(dummy_model_fn, 'gs://dummy_bucket')
        ex = Experiment(estimator, train_input_fn=None, eval_input_fn=None)
        self.assertEquals(ex._eval_every_n_steps, 1)

        # The default value when model_dir is not on GCS is 1
        estimator = Estimator(dummy_model_fn, '/tmp/dummy')
        ex = Experiment(estimator, train_input_fn=None, eval_input_fn=None)
        self.assertEquals(ex._eval_every_n_steps, 1)

        # Make sure default not used when explicitly set
        estimator = Estimator(dummy_model_fn, 'gs://dummy_bucket')
        ex = Experiment(
            estimator,
            eval_every_n_steps=123,
            train_input_fn=None,
            eval_input_fn=None)
        self.assertEquals(ex._eval_every_n_steps, 123)

        # Make sure default not used when explicitly set as 0
        estimator = Estimator(dummy_model_fn, 'gs://dummy_bucket')
        ex = Experiment(
            estimator,

def test_features_labels_mode(self):
        given_features = {'test-features': [[1], [1]]}
        given_labels = {'test-labels': [[1], [1]]}

        def _input_fn():
            return given_features, given_labels

        def _model_fn(features, labels, mode):
            self.features, self.labels, self.mode = features, labels, mode
            return EstimatorSpec(
                mode=mode,
                loss=constant_op.constant(0.),
                train_op=constant_op.constant(0.),
                predictions=constant_op.constant([[0.]]))

        est = Estimator(model_fn=_model_fn)
        est.train(_input_fn, steps=1)
        est.evaluate(_input_fn, steps=1)
        self.assertEqual(given_features, self.features)
        self.assertEqual(given_labels, self.labels)
        self.assertTrue(Modes.is_eval(self.mode))

def test_hooks_should_be_session_run_hook(self):
        est = Estimator(model_fn=model_fn_global_step_incrementer)
        est.train(dummy_input_fn, steps=1)
        with self.assertRaisesRegexp(TypeError, 'must be a SessionRunHook'):
            est.evaluate(dummy_input_fn, steps=5, hooks=['NotAHook'])

def test_evaluate_from_checkpoint(self):
        params = {
            'metric_name': 'metric',
            'metric_value': 2.}
        est1 = Estimator(
            model_fn=_model_fn_with_eval_metric_ops,
            params=params)
        est1.train(dummy_input_fn, steps=5)
        est2 = Estimator(model_fn=_model_fn_with_eval_metric_ops, params=params)
        scores = est2.evaluate(dummy_input_fn, steps=1,
                               checkpoint_path=saver.latest_checkpoint(est1.model_dir))
        self.assertEqual(5, scores['global_step'])

def test_batch_size_mismatch(self):
        def _model_fn(features, labels, mode):
            _, _ = features, labels
            return EstimatorSpec(
                mode,
                loss=constant_op.constant(0.),
                train_op=constant_op.constant(0.),
                predictions={
                    'y1': constant_op.constant([[10.]]),
                    'y2': constant_op.constant([[12.], [13]])
                })

        est = Estimator(model_fn=_model_fn)
        est.train(dummy_input_fn, steps=1)
        with self.assertRaisesRegexp(ValueError,
                                     'Batch length of predictions should be same'):
            next(est.predict(dummy_input_fn))

def test_same_model_dir_in_constructor_and_run_config(self):
        class FakeConfig(RunConfig):
            @property
            def model_dir(self):
                return _TMP_DIR

        est = Estimator(model_fn=self.get_dummy_model_fn(), config=FakeConfig(), model_dir=_TMP_DIR)
        self.assertEqual(_TMP_DIR, est.config.model_dir)
        self.assertEqual(_TMP_DIR, est.model_dir)

def graph_fn(mode, features):
    return plx.layers.SingleUnit(mode)(features['X'])


def model_fn(features, labels, mode):
    model = plx.models.Regressor(
        mode, graph_fn=graph_fn, loss_config=plx.configs.LossConfig(module='mean_squared_error'),
        optimizer_config=plx.configs.OptimizerConfig(module='sgd', learning_rate=0.009),
        eval_metrics_config=[],
        summaries='all', name='regressor')
    return model(features, labels)


estimator = plx.estimators.Estimator(model_fn=model_fn, model_dir="/tmp/polyaxon_logs/linear")

estimator.train(input_fn=numpy_input_fn(
    {'X': X}, y, shuffle=False, num_epochs=10000, batch_size=len(X)))

print([x['results'] for x in estimator.predict(input_fn=numpy_input_fn({'X': X_val},
                                                                       shuffle=False))])
print(y_val)

Recognition, 2015.
        * Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. "Gradient-based
          learning applied to document recognition." Proceedings of the IEEE,
          86(11):2278-2324, November 1998.
    Links:
        * [Deep Residual Network](http://arxiv.org/pdf/1512.03385.pdf)
        * [MNIST Dataset](http://yann.lecun.com/exdb/mnist/)

    """
    dataset_dir = '../data/mnist'
    plx.datasets.mnist.prepare(dataset_dir)
    train_input_fn, eval_input_fn = plx.datasets.mnist.create_input_fn(dataset_dir)

    run_config = plx.configs.RunConfig(save_checkpoints_steps=100)
    experiment = plx.experiments.Experiment(
        estimator=plx.estimators.Estimator(model_fn=model_fn, model_dir=output_dir,
                                           config=run_config),
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        train_steps=1000,
        eval_steps=10,
        eval_every_n_steps=5)

    return experiment

def experiment_fn(output_dir):
    """Creates an experiment using cnn for CIFAR-10 dataset classification task.

    References:
        * Learning Multiple Layers of Features from Tiny Images, A. Krizhevsky, 2009.

    Links:
        * [CIFAR-10 Dataset](https://www.cs.toronto.edu/~kriz/cifar.html)
    """
    dataset_dir = '../data/cifar10'
    plx.datasets.cifar10.prepare(dataset_dir)
    train_input_fn, eval_input_fn = plx.datasets.cifar10.create_input_fn(dataset_dir)

    experiment = plx.experiments.Experiment(
        estimator=plx.estimators.Estimator(model_fn=model_fn, model_dir=output_dir),
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        train_steps=1000,
        eval_steps=10)

    return experiment