def main(do_eval=True):
    args = parse_args()
    model_dir = get_model_dir(args)
    data = load_datasets(args.data_dir, normalize=True)
    classifier = tf.estimator.Estimator(model_fn, model_dir=model_dir)
    from kungfu.tensorflow.experimental.hook import ElasticHook
    hooks = [ElasticHook(args.batch_size, args.num_epochs, MNIST_DATA_SIZE)]
    classifier.train(input_fn(data.train, args.batch_size, epochs=args.num_epochs), hooks=hooks)
    if (not do_eval):
        import time
        time.sleep(1)
        return
    results = classifier.evaluate(input_fn(data.test, args.batch_size, shuffle=False), hooks=[], steps=1)
    print(('results: %s' % (results,)))
def main(do_eval=True):
    args = parse_args()
    model_dir = get_model_dir(args)
    data = load_datasets(args.data_dir, normalize=True)
    classifier = tf.estimator.Estimator(model_fn, model_dir=model_dir)
    from kungfu.tensorflow.experimental.hook import ElasticHook
    hooks = [ElasticHook(args.batch_size, args.num_epochs, MNIST_DATA_SIZE)]
    classifier.train(input_fn(data.train, args.batch_size, epochs=args.num_epochs), hooks=hooks)
    if (not do_eval):
        import time
        time.sleep(1)
        return
    results = classifier.evaluate(input_fn(data.test, args.batch_size, shuffle=False), hooks=[], steps=1)
    print(('results: %s' % (results,)))