We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tune] PBT (Memory Checkpointing) using TF-Saver is broken
Another issue on the same topic is that [TensorFlow examples](https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_ray_hyperband.py) does not work correctly for any training scheduler which requires checkpointing. The issue is related (as far as my understanding goes) to the fact that after running trainable.restore(...) ray deletes the temporary directory with the TensorFlow checkpoint. The TensorFlow method self.saver.restore(self.sess, path) is just a declaration of restoring process, and it does not restore model in the place where declared. The restoration of the model variables happens when session.run(...) is performed, and this is the moment when variables are loaded from saved checkpoint to the tf.Graph().  Since checkpoint directory is deleted (in trainable.restore(...) ), the actual process fails.

Why ray.tune delete checkpoint directory in trainable.restore(...)?

_Originally posted by @agniszczotka in https://github.com/ray-project/ray/issues/2856#issuecomment-459011993_
[tune] Checkpointing with tensorflow no longer works
<!--
General questions should be asked on the mailing list ray-dev@googlegroups.com.
Questions about how to use Ray should be asked on
[StackOverflow](https://stackoverflow.com/questions/tagged/ray).

Before submitting an issue, please fill out the following form.
-->

### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **Ray installed from (source or binary)**: binary
- **Ray version**: 0.6.2
- **Python version**: 3.6.7
- **Exact command to reproduce**:
https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_ray_hyperband.py


<!--
You can obtain the Ray version with

python -c "import ray; print(ray.__version__)"
-->

### Describe the problem

Saving and restoring only works with single files in the newest version for ray. Tensorflow stores multiple files during checkpointing. Thus, the tensorflow example for storing and saving is broken.

```
    def _save(self, checkpoint_dir):
        return self.saver.save(
            self.sess, checkpoint_dir + "/save", global_step=self.iterations)

    def _restore(self, path):
        return self.saver.restore(self.sess, path)
```
https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_ray_hyperband.py

### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->
ValueError: The returned checkpoint path does not exist: ray_results/2019-02-06_15-56-40yekao5r5/checkpoint_0/save


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `python/ray/tune/examples/tune_mnist_ray_hyperband.py`
Content:
```
1 #!/usr/bin/env python
2 #
3 # Copyright 2015 The TensorFlow Authors. All Rights Reserved.
4 #
5 # Licensed under the Apache License, Version 2.0 (the "License");
6 # you may not use this file except in compliance with the License.
7 # You may obtain a copy of the License at
8 #
9 #    http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing, software
12 # distributed under the License is distributed on an "AS IS" BASIS,
13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 # See the License for the specific language governing permissions and
15 # limitations under the License.
16 # ==============================================================================
17 """A deep MNIST classifier using convolutional layers.
18 See extensive documentation at
19 https://www.tensorflow.org/get_started/mnist/pros
20 """
21 # Disable linter warnings to maintain consistency with tutorial.
22 # pylint: disable=invalid-name
23 # pylint: disable=g-bad-import-order
24 
25 from __future__ import absolute_import
26 from __future__ import division
27 from __future__ import print_function
28 
29 import argparse
30 import time
31 
32 import ray
33 from ray.tune import grid_search, run_experiments, register_trainable, \
34     Trainable, sample_from
35 from ray.tune.schedulers import HyperBandScheduler
36 from tensorflow.examples.tutorials.mnist import input_data
37 
38 import tensorflow as tf
39 import numpy as np
40 
41 activation_fn = None  # e.g. tf.nn.relu
42 
43 
44 def setupCNN(x):
45     """setupCNN builds the graph for a deep net for classifying digits.
46     Args:
47         x: an input tensor with the dimensions (N_examples, 784), where 784 is
48         the number of pixels in a standard MNIST image.
49     Returns:
50         A tuple (y, keep_prob). y is a tensor of shape (N_examples, 10), with
51         values equal to the logits of classifying the digit into one of 10
52         classes (the digits 0-9). keep_prob is a scalar placeholder for the
53         probability of dropout.
54     """
55     # Reshape to use within a convolutional neural net.
56     # Last dimension is for "features" - there is only one here, since images
57     # are grayscale -- it would be 3 for an RGB image, 4 for RGBA, etc.
58     with tf.name_scope('reshape'):
59         x_image = tf.reshape(x, [-1, 28, 28, 1])
60 
61     # First convolutional layer - maps one grayscale image to 32 feature maps.
62     with tf.name_scope('conv1'):
63         W_conv1 = weight_variable([5, 5, 1, 32])
64         b_conv1 = bias_variable([32])
65         h_conv1 = activation_fn(conv2d(x_image, W_conv1) + b_conv1)
66 
67     # Pooling layer - downsamples by 2X.
68     with tf.name_scope('pool1'):
69         h_pool1 = max_pool_2x2(h_conv1)
70 
71     # Second convolutional layer -- maps 32 feature maps to 64.
72     with tf.name_scope('conv2'):
73         W_conv2 = weight_variable([5, 5, 32, 64])
74         b_conv2 = bias_variable([64])
75         h_conv2 = activation_fn(conv2d(h_pool1, W_conv2) + b_conv2)
76 
77     # Second pooling layer.
78     with tf.name_scope('pool2'):
79         h_pool2 = max_pool_2x2(h_conv2)
80 
81     # Fully connected layer 1 -- after 2 round of downsampling, our 28x28 image
82     # is down to 7x7x64 feature maps -- maps this to 1024 features.
83     with tf.name_scope('fc1'):
84         W_fc1 = weight_variable([7 * 7 * 64, 1024])
85         b_fc1 = bias_variable([1024])
86 
87         h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])
88         h_fc1 = activation_fn(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)
89 
90     # Dropout - controls the complexity of the model, prevents co-adaptation of
91     # features.
92     with tf.name_scope('dropout'):
93         keep_prob = tf.placeholder(tf.float32)
94         h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)
95 
96     # Map the 1024 features to 10 classes, one for each digit
97     with tf.name_scope('fc2'):
98         W_fc2 = weight_variable([1024, 10])
99         b_fc2 = bias_variable([10])
100 
101         y_conv = tf.matmul(h_fc1_drop, W_fc2) + b_fc2
102     return y_conv, keep_prob
103 
104 
105 def conv2d(x, W):
106     """conv2d returns a 2d convolution layer with full stride."""
107     return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')
108 
109 
110 def max_pool_2x2(x):
111     """max_pool_2x2 downsamples a feature map by 2X."""
112     return tf.nn.max_pool(
113         x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
114 
115 
116 def weight_variable(shape):
117     """weight_variable generates a weight variable of a given shape."""
118     initial = tf.truncated_normal(shape, stddev=0.1)
119     return tf.Variable(initial)
120 
121 
122 def bias_variable(shape):
123     """bias_variable generates a bias variable of a given shape."""
124     initial = tf.constant(0.1, shape=shape)
125     return tf.Variable(initial)
126 
127 
128 class TrainMNIST(Trainable):
129     """Example MNIST trainable."""
130 
131     def _setup(self, config):
132         global activation_fn
133 
134         self.timestep = 0
135 
136         # Import data
137         for _ in range(10):
138             try:
139                 self.mnist = input_data.read_data_sets(
140                     "/tmp/mnist_ray_demo", one_hot=True)
141                 break
142             except Exception as e:
143                 print("Error loading data, retrying", e)
144                 time.sleep(5)
145 
146         assert self.mnist
147 
148         self.x = tf.placeholder(tf.float32, [None, 784])
149         self.y_ = tf.placeholder(tf.float32, [None, 10])
150 
151         activation_fn = getattr(tf.nn, config['activation'])
152 
153         # Build the graph for the deep net
154         y_conv, self.keep_prob = setupCNN(self.x)
155 
156         with tf.name_scope('loss'):
157             cross_entropy = tf.nn.softmax_cross_entropy_with_logits(
158                 labels=self.y_, logits=y_conv)
159         cross_entropy = tf.reduce_mean(cross_entropy)
160 
161         with tf.name_scope('adam_optimizer'):
162             train_step = tf.train.AdamOptimizer(
163                 config['learning_rate']).minimize(cross_entropy)
164 
165         self.train_step = train_step
166 
167         with tf.name_scope('accuracy'):
168             correct_prediction = tf.equal(
169                 tf.argmax(y_conv, 1), tf.argmax(self.y_, 1))
170             correct_prediction = tf.cast(correct_prediction, tf.float32)
171         self.accuracy = tf.reduce_mean(correct_prediction)
172 
173         self.sess = tf.Session()
174         self.sess.run(tf.global_variables_initializer())
175         self.iterations = 0
176         self.saver = tf.train.Saver()
177 
178     def _train(self):
179         for i in range(10):
180             batch = self.mnist.train.next_batch(50)
181             self.sess.run(
182                 self.train_step,
183                 feed_dict={
184                     self.x: batch[0],
185                     self.y_: batch[1],
186                     self.keep_prob: 0.5
187                 })
188 
189         batch = self.mnist.train.next_batch(50)
190         train_accuracy = self.sess.run(
191             self.accuracy,
192             feed_dict={
193                 self.x: batch[0],
194                 self.y_: batch[1],
195                 self.keep_prob: 1.0
196             })
197 
198         self.iterations += 1
199         return {"mean_accuracy": train_accuracy}
200 
201     def _save(self, checkpoint_dir):
202         return self.saver.save(
203             self.sess, checkpoint_dir + "/save", global_step=self.iterations)
204 
205     def _restore(self, path):
206         return self.saver.restore(self.sess, path)
207 
208 
209 # !!! Example of using the ray.tune Python API !!!
210 if __name__ == '__main__':
211     parser = argparse.ArgumentParser()
212     parser.add_argument(
213         '--smoke-test', action='store_true', help='Finish quickly for testing')
214     args, _ = parser.parse_known_args()
215 
216     register_trainable("my_class", TrainMNIST)
217     mnist_spec = {
218         'run': 'my_class',
219         'stop': {
220             'mean_accuracy': 0.99,
221             'time_total_s': 600,
222         },
223         'config': {
224             'learning_rate': sample_from(
225                 lambda spec: 10**np.random.uniform(-5, -3)),
226             'activation': grid_search(['relu', 'elu', 'tanh']),
227         },
228         "num_samples": 10,
229     }
230 
231     if args.smoke_test:
232         mnist_spec['stop']['training_iteration'] = 2
233         mnist_spec['num_samples'] = 2
234 
235     ray.init()
236     hyperband = HyperBandScheduler(
237         time_attr="training_iteration", reward_attr="mean_accuracy", max_t=10)
238 
239     run_experiments({'mnist_hyperband_test': mnist_spec}, scheduler=hyperband)
240
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/python/ray/tune/examples/tune_mnist_ray_hyperband.py b/python/ray/tune/examples/tune_mnist_ray_hyperband.py
--- a/python/ray/tune/examples/tune_mnist_ray_hyperband.py
+++ b/python/ray/tune/examples/tune_mnist_ray_hyperband.py
@@ -199,11 +199,13 @@
         return {"mean_accuracy": train_accuracy}
 
     def _save(self, checkpoint_dir):
-        return self.saver.save(
+        prefix = self.saver.save(
             self.sess, checkpoint_dir + "/save", global_step=self.iterations)
+        return {"prefix": prefix}
 
-    def _restore(self, path):
-        return self.saver.restore(self.sess, path)
+    def _restore(self, ckpt_data):
+        prefix = ckpt_data["prefix"]
+        return self.saver.restore(self.sess, prefix)
 
 
 # !!! Example of using the ray.tune Python API !!!
@@ -229,7 +231,7 @@
     }
 
     if args.smoke_test:
-        mnist_spec['stop']['training_iteration'] = 2
+        mnist_spec['stop']['training_iteration'] = 20
         mnist_spec['num_samples'] = 2
 
     ray.init()

{"golden_diff": "diff --git a/python/ray/tune/examples/tune_mnist_ray_hyperband.py b/python/ray/tune/examples/tune_mnist_ray_hyperband.py\n--- a/python/ray/tune/examples/tune_mnist_ray_hyperband.py\n+++ b/python/ray/tune/examples/tune_mnist_ray_hyperband.py\n@@ -199,11 +199,13 @@\n         return {\"mean_accuracy\": train_accuracy}\n \n     def _save(self, checkpoint_dir):\n-        return self.saver.save(\n+        prefix = self.saver.save(\n             self.sess, checkpoint_dir + \"/save\", global_step=self.iterations)\n+        return {\"prefix\": prefix}\n \n-    def _restore(self, path):\n-        return self.saver.restore(self.sess, path)\n+    def _restore(self, ckpt_data):\n+        prefix = ckpt_data[\"prefix\"]\n+        return self.saver.restore(self.sess, prefix)\n \n \n # !!! Example of using the ray.tune Python API !!!\n@@ -229,7 +231,7 @@\n     }\n \n     if args.smoke_test:\n-        mnist_spec['stop']['training_iteration'] = 2\n+        mnist_spec['stop']['training_iteration'] = 20\n         mnist_spec['num_samples'] = 2\n \n     ray.init()\n", "issue": "[tune] PBT (Memory Checkpointing) using TF-Saver is broken\nAnother issue on the same topic is that [TensorFlow examples](https://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_ray_hyperband.py) does not work correctly for any training scheduler which requires checkpointing. The issue is related (as far as my understanding goes) to the fact that after running trainable.restore(...) ray deletes the temporary directory with the TensorFlow checkpoint. The TensorFlow method self.saver.restore(self.sess, path) is just a declaration of restoring process, and it does not restore model in the place where declared. The restoration of the model variables happens when session.run(...) is performed, and this is the moment when variables are loaded from saved checkpoint to the tf.Graph().  Since checkpoint directory is deleted (in trainable.restore(...) ), the actual process fails.\r\n\r\nWhy ray.tune delete checkpoint directory in trainable.restore(...)?\r\n\r\n_Originally posted by @agniszczotka in https://github.com/ray-project/ray/issues/2856#issuecomment-459011993_\n[tune] Checkpointing with tensorflow no longer works\n<!--\r\nGeneral questions should be asked on the mailing list ray-dev@googlegroups.com.\r\nQuestions about how to use Ray should be asked on\r\n[StackOverflow](https://stackoverflow.com/questions/tagged/ray).\r\n\r\nBefore submitting an issue, please fill out the following form.\r\n-->\r\n\r\n### System information\r\n- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:\r\n- **Ray installed from (source or binary)**: binary\r\n- **Ray version**: 0.6.2\r\n- **Python version**: 3.6.7\r\n- **Exact command to reproduce**:\r\nhttps://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_ray_hyperband.py\r\n\r\n\r\n<!--\r\nYou can obtain the Ray version with\r\n\r\npython -c \"import ray; print(ray.__version__)\"\r\n-->\r\n\r\n### Describe the problem\r\n\r\nSaving and restoring only works with single files in the newest version for ray. Tensorflow stores multiple files during checkpointing. Thus, the tensorflow example for storing and saving is broken.\r\n\r\n```\r\n    def _save(self, checkpoint_dir):\r\n        return self.saver.save(\r\n            self.sess, checkpoint_dir + \"/save\", global_step=self.iterations)\r\n\r\n    def _restore(self, path):\r\n        return self.saver.restore(self.sess, path)\r\n```\r\nhttps://github.com/ray-project/ray/blob/master/python/ray/tune/examples/tune_mnist_ray_hyperband.py\r\n\r\n### Source code / logs\r\n<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->\r\nValueError: The returned checkpoint path does not exist: ray_results/2019-02-06_15-56-40yekao5r5/checkpoint_0/save\r\n\n", "before_files": [{"content": "#!/usr/bin/env python\n#\n# Copyright 2015 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"A deep MNIST classifier using convolutional layers.\nSee extensive documentation at\nhttps://www.tensorflow.org/get_started/mnist/pros\n\"\"\"\n# Disable linter warnings to maintain consistency with tutorial.\n# pylint: disable=invalid-name\n# pylint: disable=g-bad-import-order\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport argparse\nimport time\n\nimport ray\nfrom ray.tune import grid_search, run_experiments, register_trainable, \\\n    Trainable, sample_from\nfrom ray.tune.schedulers import HyperBandScheduler\nfrom tensorflow.examples.tutorials.mnist import input_data\n\nimport tensorflow as tf\nimport numpy as np\n\nactivation_fn = None  # e.g. tf.nn.relu\n\n\ndef setupCNN(x):\n    \"\"\"setupCNN builds the graph for a deep net for classifying digits.\n    Args:\n        x: an input tensor with the dimensions (N_examples, 784), where 784 is\n        the number of pixels in a standard MNIST image.\n    Returns:\n        A tuple (y, keep_prob). y is a tensor of shape (N_examples, 10), with\n        values equal to the logits of classifying the digit into one of 10\n        classes (the digits 0-9). keep_prob is a scalar placeholder for the\n        probability of dropout.\n    \"\"\"\n    # Reshape to use within a convolutional neural net.\n    # Last dimension is for \"features\" - there is only one here, since images\n    # are grayscale -- it would be 3 for an RGB image, 4 for RGBA, etc.\n    with tf.name_scope('reshape'):\n        x_image = tf.reshape(x, [-1, 28, 28, 1])\n\n    # First convolutional layer - maps one grayscale image to 32 feature maps.\n    with tf.name_scope('conv1'):\n        W_conv1 = weight_variable([5, 5, 1, 32])\n        b_conv1 = bias_variable([32])\n        h_conv1 = activation_fn(conv2d(x_image, W_conv1) + b_conv1)\n\n    # Pooling layer - downsamples by 2X.\n    with tf.name_scope('pool1'):\n        h_pool1 = max_pool_2x2(h_conv1)\n\n    # Second convolutional layer -- maps 32 feature maps to 64.\n    with tf.name_scope('conv2'):\n        W_conv2 = weight_variable([5, 5, 32, 64])\n        b_conv2 = bias_variable([64])\n        h_conv2 = activation_fn(conv2d(h_pool1, W_conv2) + b_conv2)\n\n    # Second pooling layer.\n    with tf.name_scope('pool2'):\n        h_pool2 = max_pool_2x2(h_conv2)\n\n    # Fully connected layer 1 -- after 2 round of downsampling, our 28x28 image\n    # is down to 7x7x64 feature maps -- maps this to 1024 features.\n    with tf.name_scope('fc1'):\n        W_fc1 = weight_variable([7 * 7 * 64, 1024])\n        b_fc1 = bias_variable([1024])\n\n        h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])\n        h_fc1 = activation_fn(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)\n\n    # Dropout - controls the complexity of the model, prevents co-adaptation of\n    # features.\n    with tf.name_scope('dropout'):\n        keep_prob = tf.placeholder(tf.float32)\n        h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)\n\n    # Map the 1024 features to 10 classes, one for each digit\n    with tf.name_scope('fc2'):\n        W_fc2 = weight_variable([1024, 10])\n        b_fc2 = bias_variable([10])\n\n        y_conv = tf.matmul(h_fc1_drop, W_fc2) + b_fc2\n    return y_conv, keep_prob\n\n\ndef conv2d(x, W):\n    \"\"\"conv2d returns a 2d convolution layer with full stride.\"\"\"\n    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')\n\n\ndef max_pool_2x2(x):\n    \"\"\"max_pool_2x2 downsamples a feature map by 2X.\"\"\"\n    return tf.nn.max_pool(\n        x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')\n\n\ndef weight_variable(shape):\n    \"\"\"weight_variable generates a weight variable of a given shape.\"\"\"\n    initial = tf.truncated_normal(shape, stddev=0.1)\n    return tf.Variable(initial)\n\n\ndef bias_variable(shape):\n    \"\"\"bias_variable generates a bias variable of a given shape.\"\"\"\n    initial = tf.constant(0.1, shape=shape)\n    return tf.Variable(initial)\n\n\nclass TrainMNIST(Trainable):\n    \"\"\"Example MNIST trainable.\"\"\"\n\n    def _setup(self, config):\n        global activation_fn\n\n        self.timestep = 0\n\n        # Import data\n        for _ in range(10):\n            try:\n                self.mnist = input_data.read_data_sets(\n                    \"/tmp/mnist_ray_demo\", one_hot=True)\n                break\n            except Exception as e:\n                print(\"Error loading data, retrying\", e)\n                time.sleep(5)\n\n        assert self.mnist\n\n        self.x = tf.placeholder(tf.float32, [None, 784])\n        self.y_ = tf.placeholder(tf.float32, [None, 10])\n\n        activation_fn = getattr(tf.nn, config['activation'])\n\n        # Build the graph for the deep net\n        y_conv, self.keep_prob = setupCNN(self.x)\n\n        with tf.name_scope('loss'):\n            cross_entropy = tf.nn.softmax_cross_entropy_with_logits(\n                labels=self.y_, logits=y_conv)\n        cross_entropy = tf.reduce_mean(cross_entropy)\n\n        with tf.name_scope('adam_optimizer'):\n            train_step = tf.train.AdamOptimizer(\n                config['learning_rate']).minimize(cross_entropy)\n\n        self.train_step = train_step\n\n        with tf.name_scope('accuracy'):\n            correct_prediction = tf.equal(\n                tf.argmax(y_conv, 1), tf.argmax(self.y_, 1))\n            correct_prediction = tf.cast(correct_prediction, tf.float32)\n        self.accuracy = tf.reduce_mean(correct_prediction)\n\n        self.sess = tf.Session()\n        self.sess.run(tf.global_variables_initializer())\n        self.iterations = 0\n        self.saver = tf.train.Saver()\n\n    def _train(self):\n        for i in range(10):\n            batch = self.mnist.train.next_batch(50)\n            self.sess.run(\n                self.train_step,\n                feed_dict={\n                    self.x: batch[0],\n                    self.y_: batch[1],\n                    self.keep_prob: 0.5\n                })\n\n        batch = self.mnist.train.next_batch(50)\n        train_accuracy = self.sess.run(\n            self.accuracy,\n            feed_dict={\n                self.x: batch[0],\n                self.y_: batch[1],\n                self.keep_prob: 1.0\n            })\n\n        self.iterations += 1\n        return {\"mean_accuracy\": train_accuracy}\n\n    def _save(self, checkpoint_dir):\n        return self.saver.save(\n            self.sess, checkpoint_dir + \"/save\", global_step=self.iterations)\n\n    def _restore(self, path):\n        return self.saver.restore(self.sess, path)\n\n\n# !!! Example of using the ray.tune Python API !!!\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\n        '--smoke-test', action='store_true', help='Finish quickly for testing')\n    args, _ = parser.parse_known_args()\n\n    register_trainable(\"my_class\", TrainMNIST)\n    mnist_spec = {\n        'run': 'my_class',\n        'stop': {\n            'mean_accuracy': 0.99,\n            'time_total_s': 600,\n        },\n        'config': {\n            'learning_rate': sample_from(\n                lambda spec: 10**np.random.uniform(-5, -3)),\n            'activation': grid_search(['relu', 'elu', 'tanh']),\n        },\n        \"num_samples\": 10,\n    }\n\n    if args.smoke_test:\n        mnist_spec['stop']['training_iteration'] = 2\n        mnist_spec['num_samples'] = 2\n\n    ray.init()\n    hyperband = HyperBandScheduler(\n        time_attr=\"training_iteration\", reward_attr=\"mean_accuracy\", max_t=10)\n\n    run_experiments({'mnist_hyperband_test': mnist_spec}, scheduler=hyperband)\n", "path": "python/ray/tune/examples/tune_mnist_ray_hyperband.py"}], "after_files": [{"content": "#!/usr/bin/env python\n#\n# Copyright 2015 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"A deep MNIST classifier using convolutional layers.\nSee extensive documentation at\nhttps://www.tensorflow.org/get_started/mnist/pros\n\"\"\"\n# Disable linter warnings to maintain consistency with tutorial.\n# pylint: disable=invalid-name\n# pylint: disable=g-bad-import-order\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport argparse\nimport time\n\nimport ray\nfrom ray.tune import grid_search, run_experiments, register_trainable, \\\n    Trainable, sample_from\nfrom ray.tune.schedulers import HyperBandScheduler\nfrom tensorflow.examples.tutorials.mnist import input_data\n\nimport tensorflow as tf\nimport numpy as np\n\nactivation_fn = None  # e.g. tf.nn.relu\n\n\ndef setupCNN(x):\n    \"\"\"setupCNN builds the graph for a deep net for classifying digits.\n    Args:\n        x: an input tensor with the dimensions (N_examples, 784), where 784 is\n        the number of pixels in a standard MNIST image.\n    Returns:\n        A tuple (y, keep_prob). y is a tensor of shape (N_examples, 10), with\n        values equal to the logits of classifying the digit into one of 10\n        classes (the digits 0-9). keep_prob is a scalar placeholder for the\n        probability of dropout.\n    \"\"\"\n    # Reshape to use within a convolutional neural net.\n    # Last dimension is for \"features\" - there is only one here, since images\n    # are grayscale -- it would be 3 for an RGB image, 4 for RGBA, etc.\n    with tf.name_scope('reshape'):\n        x_image = tf.reshape(x, [-1, 28, 28, 1])\n\n    # First convolutional layer - maps one grayscale image to 32 feature maps.\n    with tf.name_scope('conv1'):\n        W_conv1 = weight_variable([5, 5, 1, 32])\n        b_conv1 = bias_variable([32])\n        h_conv1 = activation_fn(conv2d(x_image, W_conv1) + b_conv1)\n\n    # Pooling layer - downsamples by 2X.\n    with tf.name_scope('pool1'):\n        h_pool1 = max_pool_2x2(h_conv1)\n\n    # Second convolutional layer -- maps 32 feature maps to 64.\n    with tf.name_scope('conv2'):\n        W_conv2 = weight_variable([5, 5, 32, 64])\n        b_conv2 = bias_variable([64])\n        h_conv2 = activation_fn(conv2d(h_pool1, W_conv2) + b_conv2)\n\n    # Second pooling layer.\n    with tf.name_scope('pool2'):\n        h_pool2 = max_pool_2x2(h_conv2)\n\n    # Fully connected layer 1 -- after 2 round of downsampling, our 28x28 image\n    # is down to 7x7x64 feature maps -- maps this to 1024 features.\n    with tf.name_scope('fc1'):\n        W_fc1 = weight_variable([7 * 7 * 64, 1024])\n        b_fc1 = bias_variable([1024])\n\n        h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])\n        h_fc1 = activation_fn(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)\n\n    # Dropout - controls the complexity of the model, prevents co-adaptation of\n    # features.\n    with tf.name_scope('dropout'):\n        keep_prob = tf.placeholder(tf.float32)\n        h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)\n\n    # Map the 1024 features to 10 classes, one for each digit\n    with tf.name_scope('fc2'):\n        W_fc2 = weight_variable([1024, 10])\n        b_fc2 = bias_variable([10])\n\n        y_conv = tf.matmul(h_fc1_drop, W_fc2) + b_fc2\n    return y_conv, keep_prob\n\n\ndef conv2d(x, W):\n    \"\"\"conv2d returns a 2d convolution layer with full stride.\"\"\"\n    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')\n\n\ndef max_pool_2x2(x):\n    \"\"\"max_pool_2x2 downsamples a feature map by 2X.\"\"\"\n    return tf.nn.max_pool(\n        x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')\n\n\ndef weight_variable(shape):\n    \"\"\"weight_variable generates a weight variable of a given shape.\"\"\"\n    initial = tf.truncated_normal(shape, stddev=0.1)\n    return tf.Variable(initial)\n\n\ndef bias_variable(shape):\n    \"\"\"bias_variable generates a bias variable of a given shape.\"\"\"\n    initial = tf.constant(0.1, shape=shape)\n    return tf.Variable(initial)\n\n\nclass TrainMNIST(Trainable):\n    \"\"\"Example MNIST trainable.\"\"\"\n\n    def _setup(self, config):\n        global activation_fn\n\n        self.timestep = 0\n\n        # Import data\n        for _ in range(10):\n            try:\n                self.mnist = input_data.read_data_sets(\n                    \"/tmp/mnist_ray_demo\", one_hot=True)\n                break\n            except Exception as e:\n                print(\"Error loading data, retrying\", e)\n                time.sleep(5)\n\n        assert self.mnist\n\n        self.x = tf.placeholder(tf.float32, [None, 784])\n        self.y_ = tf.placeholder(tf.float32, [None, 10])\n\n        activation_fn = getattr(tf.nn, config['activation'])\n\n        # Build the graph for the deep net\n        y_conv, self.keep_prob = setupCNN(self.x)\n\n        with tf.name_scope('loss'):\n            cross_entropy = tf.nn.softmax_cross_entropy_with_logits(\n                labels=self.y_, logits=y_conv)\n        cross_entropy = tf.reduce_mean(cross_entropy)\n\n        with tf.name_scope('adam_optimizer'):\n            train_step = tf.train.AdamOptimizer(\n                config['learning_rate']).minimize(cross_entropy)\n\n        self.train_step = train_step\n\n        with tf.name_scope('accuracy'):\n            correct_prediction = tf.equal(\n                tf.argmax(y_conv, 1), tf.argmax(self.y_, 1))\n            correct_prediction = tf.cast(correct_prediction, tf.float32)\n        self.accuracy = tf.reduce_mean(correct_prediction)\n\n        self.sess = tf.Session()\n        self.sess.run(tf.global_variables_initializer())\n        self.iterations = 0\n        self.saver = tf.train.Saver()\n\n    def _train(self):\n        for i in range(10):\n            batch = self.mnist.train.next_batch(50)\n            self.sess.run(\n                self.train_step,\n                feed_dict={\n                    self.x: batch[0],\n                    self.y_: batch[1],\n                    self.keep_prob: 0.5\n                })\n\n        batch = self.mnist.train.next_batch(50)\n        train_accuracy = self.sess.run(\n            self.accuracy,\n            feed_dict={\n                self.x: batch[0],\n                self.y_: batch[1],\n                self.keep_prob: 1.0\n            })\n\n        self.iterations += 1\n        return {\"mean_accuracy\": train_accuracy}\n\n    def _save(self, checkpoint_dir):\n        prefix = self.saver.save(\n            self.sess, checkpoint_dir + \"/save\", global_step=self.iterations)\n        return {\"prefix\": prefix}\n\n    def _restore(self, ckpt_data):\n        prefix = ckpt_data[\"prefix\"]\n        return self.saver.restore(self.sess, prefix)\n\n\n# !!! Example of using the ray.tune Python API !!!\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\n        '--smoke-test', action='store_true', help='Finish quickly for testing')\n    args, _ = parser.parse_known_args()\n\n    register_trainable(\"my_class\", TrainMNIST)\n    mnist_spec = {\n        'run': 'my_class',\n        'stop': {\n            'mean_accuracy': 0.99,\n            'time_total_s': 600,\n        },\n        'config': {\n            'learning_rate': sample_from(\n                lambda spec: 10**np.random.uniform(-5, -3)),\n            'activation': grid_search(['relu', 'elu', 'tanh']),\n        },\n        \"num_samples\": 10,\n    }\n\n    if args.smoke_test:\n        mnist_spec['stop']['training_iteration'] = 20\n        mnist_spec['num_samples'] = 2\n\n    ray.init()\n    hyperband = HyperBandScheduler(\n        time_attr=\"training_iteration\", reward_attr=\"mean_accuracy\", max_t=10)\n\n    run_experiments({'mnist_hyperband_test': mnist_spec}, scheduler=hyperband)\n", "path": "python/ray/tune/examples/tune_mnist_ray_hyperband.py"}]}