We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
About ENAS_Scheduler
While using ENAS_Scheduler
Maybe：
```
tbar.set_description('avg reward: {}'.format(self.baseline))
```
is better than 

```
tbar.set_description('avg reward: {:.2f}'.format(self.baseline))
```
in run in enas_scheduler.py, because when self.baseline=None, it will report error

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `autogluon/contrib/enas/enas_scheduler.py`
Content:
```
1 import os
2 import pickle
3 import logging
4 from collections import OrderedDict
5 from multiprocessing.pool import ThreadPool
6 
7 import mxnet as mx
8 
9 from ...searcher import RLSearcher
10 from ...scheduler.resource import get_gpu_count, get_cpu_count
11 from ...task.image_classification.dataset import get_built_in_dataset
12 from ...task.image_classification.utils import *
13 from ...utils import (mkdir, save, load, update_params, collect_params, DataLoader, tqdm, in_ipynb)
14 from .enas_utils import *
15 
16 __all__ = ['ENAS_Scheduler']
17 
18 logger = logging.getLogger(__name__)
19 
20 IMAGENET_TRAINING_SAMPLES = 1281167
21 
22 class ENAS_Scheduler(object):
23     """ENAS Scheduler, which automatically creates LSTM controller based on the search spaces.
24     """
25     def __init__(self, supernet, train_set='imagenet', val_set=None,
26                  train_fn=default_train_fn, eval_fn=default_val_fn,
27                  train_args={}, val_args={}, reward_fn= default_reward_fn,
28                  num_gpus=0, num_cpus=4,
29                  batch_size=256, epochs=120, warmup_epochs=5,
30                  controller_lr=1e-3, controller_type='lstm',
31                  controller_batch_size=10, ema_baseline_decay=0.95,
32                  update_arch_frequency=20, checkname='./enas/checkpoint.ag',
33                  plot_frequency=0, **kwargs):
34         num_cpus = get_cpu_count() if num_cpus > get_cpu_count() else num_cpus
35         num_gpus = get_gpu_count() if num_gpus > get_gpu_count() else num_gpus
36         self.supernet = supernet
37         self.train_fn = train_fn
38         self.eval_fn = eval_fn
39         self.reward_fn = reward_fn
40         self.checkname = checkname
41         self.plot_frequency = plot_frequency
42         self.epochs = epochs
43         self.warmup_epochs = warmup_epochs
44         self.controller_batch_size = controller_batch_size
45         kwspaces = self.supernet.kwspaces
46 
47         self.initialize_miscs(train_set, val_set, batch_size, num_cpus, num_gpus,
48                               train_args, val_args)
49 
50         # create RL searcher/controller
51         self.baseline = None
52         self.ema_decay = ema_baseline_decay
53         self.searcher = RLSearcher(kwspaces, controller_type=controller_type,
54                                    prefetch=4, num_workers=4)
55         # controller setup
56         self.controller = self.searcher.controller
57         self.controller_optimizer = mx.gluon.Trainer(
58                 self.controller.collect_params(), 'adam',
59                 optimizer_params={'learning_rate': controller_lr})
60         self.update_arch_frequency = update_arch_frequency
61         self.val_acc = 0
62         # async controller sample
63         self._worker_pool = ThreadPool(2)
64         self._data_buffer = {}
65         self._rcvd_idx = 0
66         self._sent_idx = 0
67         self._timeout = 20
68         # logging history
69         self.training_history = []
70         self._prefetch_controller()
71 
72     def initialize_miscs(self, train_set, val_set, batch_size, num_cpus, num_gpus,
73                          train_args, val_args):
74         """Initialize framework related miscs, such as train/val data and train/val
75         function arguments.
76         """
77         ctx = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu(0)]
78         self.supernet.collect_params().reset_ctx(ctx)
79         self.supernet.hybridize()
80         dataset_name = train_set
81 
82         if isinstance(train_set, str):
83             train_set = get_built_in_dataset(dataset_name, train=True, batch_size=batch_size,
84                                              num_workers=num_cpus, shuffle=True)
85             val_set = get_built_in_dataset(dataset_name, train=False, batch_size=batch_size,
86                                            num_workers=num_cpus, shuffle=True)
87         if isinstance(train_set, gluon.data.Dataset):
88             self.train_data = DataLoader(
89                     train_set, batch_size=batch_size, shuffle=True,
90                     last_batch="discard", num_workers=num_cpus)
91             # very important, make shuffle for training contoller
92             self.val_data = DataLoader(
93                     val_set, batch_size=batch_size, shuffle=True,
94                     num_workers=num_cpus, prefetch=0, sample_times=self.controller_batch_size)
95         else:
96             self.train_data = train_set
97             self.val_data = val_set
98         iters_per_epoch = len(self.train_data) if hasattr(self.train_data, '__len__') else \
99                 IMAGENET_TRAINING_SAMPLES // batch_size
100         self.train_args = init_default_train_args(batch_size, self.supernet, self.epochs, iters_per_epoch) \
101                 if len(train_args) == 0 else train_args
102         self.val_args = val_args
103         self.val_args['ctx'] = ctx
104         self.val_args['batch_fn'] = imagenet_batch_fn if dataset_name == 'imagenet' else default_batch_fn
105         self.train_args['ctx'] = ctx
106         self.train_args['batch_fn'] = imagenet_batch_fn if dataset_name == 'imagenet' else default_batch_fn
107         self.ctx = ctx
108 
109     def run(self):
110         tq = tqdm(range(self.epochs))
111         for epoch in tq:
112             # for recordio data
113             if hasattr(self.train_data, 'reset'): self.train_data.reset()
114             tbar = tqdm(self.train_data)
115             idx = 0
116             for batch in tbar:
117                 # sample network configuration
118                 config = self.controller.pre_sample()[0]
119                 self.supernet.sample(**config)
120                 # self.train_fn(self.supernet, batch, **self.train_args)
121                 self.train_fn(epoch, self.epochs, self.supernet, batch, **self.train_args)
122                 mx.nd.waitall()
123                 if epoch >= self.warmup_epochs and (idx % self.update_arch_frequency) == 0:
124                     self.train_controller()
125                 if self.plot_frequency > 0 and idx % self.plot_frequency == 0 and in_ipynb():
126                     graph = self.supernet.graph
127                     graph.attr(rankdir='LR', size='8,3')
128                     tbar.set_svg(graph._repr_svg_())
129                 tbar.set_description('avg reward: {:.2f}'.format(self.baseline))
130                 idx += 1
131             self.validation()
132             self.save()
133             tq.set_description('epoch {}, val_acc: {:.2f}, avg reward: {:.2f}' \
134                         .format(epoch, self.val_acc, self.baseline))
135 
136     def validation(self):
137         if hasattr(self.val_data, 'reset'): self.val_data.reset()
138         # data iter, avoid memory leak
139         it = iter(self.val_data)
140         if hasattr(it, 'reset_sample_times'): it.reset_sample_times()
141         tbar = tqdm(it)
142         # update network arc
143         config = self.controller.inference()
144         self.supernet.sample(**config)
145         metric = mx.metric.Accuracy()
146         for batch in tbar:
147             self.eval_fn(self.supernet, batch, metric=metric, **self.val_args)
148             reward = metric.get()[1]
149             tbar.set_description('Val Acc: {}'.format(reward))
150 
151         self.val_acc = reward
152         self.training_history.append(reward)
153 
154     def _sample_controller(self):
155         assert self._rcvd_idx < self._sent_idx, "rcvd_idx must be smaller than sent_idx"
156         try:
157             ret = self._data_buffer.pop(self._rcvd_idx)
158             self._rcvd_idx += 1
159             return  ret.get(timeout=self._timeout)
160         except Exception:
161             self._worker_pool.terminate()
162             raise
163 
164     def _prefetch_controller(self):
165         async_ret = self._worker_pool.apply_async(self._async_sample, ())
166         self._data_buffer[self._sent_idx] = async_ret
167         self._sent_idx += 1
168 
169     def _async_sample(self):
170         with mx.autograd.record():
171             # sample controller_batch_size number of configurations
172             configs, log_probs, entropies = self.controller.sample(batch_size=self.controller_batch_size,
173                                                                    with_details=True)
174         return configs, log_probs, entropies
175 
176     def train_controller(self):
177         """Run multiple number of trials
178         """
179         decay = self.ema_decay
180         if hasattr(self.val_data, 'reset'): self.val_data.reset()
181         # update 
182         metric = mx.metric.Accuracy()
183         with mx.autograd.record():
184             # sample controller_batch_size number of configurations
185             configs, log_probs, entropies = self._sample_controller()
186             for i, batch in enumerate(self.val_data):
187                 if i >= self.controller_batch_size: break
188                 self.supernet.sample(**configs[i])
189                 # schedule the training tasks and gather the reward
190                 metric.reset()
191                 self.eval_fn(self.supernet, batch, metric=metric, **self.val_args)
192                 reward = metric.get()[1]
193                 reward = self.reward_fn(reward, self.supernet)
194                 self.baseline = reward if not self.baseline else self.baseline
195                 # substract baseline
196                 avg_rewards = mx.nd.array([reward - self.baseline],
197                                           ctx=self.controller.context)
198                 # EMA baseline
199                 self.baseline = decay * self.baseline + (1 - decay) * reward
200                 # negative policy gradient
201                 log_prob = log_probs[i]
202                 log_prob = log_prob.sum()
203                 loss = - log_prob * avg_rewards
204                 loss = loss.sum()
205 
206         # update
207         loss.backward()
208         self.controller_optimizer.step(self.controller_batch_size)
209         self._prefetch_controller()
210 
211     def load(self, checkname=None):
212         checkname = checkname if checkname else self.checkname
213         state_dict = load(checkname)
214         self.load_state_dict(state_dict)
215 
216     def save(self, checkname=None):
217         checkname = checkname if checkname else self.checkname
218         mkdir(os.path.dirname(checkname))
219         save(self.state_dict(), checkname)
220 
221     def state_dict(self, destination=None):
222         if destination is None:
223             destination = OrderedDict()
224             destination._metadata = OrderedDict()
225         destination['supernet_params'] = collect_params(self.supernet)
226         destination['controller_params'] = collect_params(self.controller)
227         destination['training_history'] = self.training_history
228         return destination
229 
230     def load_state_dict(self, state_dict):
231         update_params(self.supernet, state_dict['supernet_params'], ctx=self.ctx)
232         update_params(self.controller, state_dict['controller_params'], ctx=self.controller.context)
233         self.training_history = state_dict['training_history']
234
```
--- 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/autogluon/contrib/enas/enas_scheduler.py b/autogluon/contrib/enas/enas_scheduler.py
--- a/autogluon/contrib/enas/enas_scheduler.py
+++ b/autogluon/contrib/enas/enas_scheduler.py
@@ -126,12 +126,15 @@
                     graph = self.supernet.graph
                     graph.attr(rankdir='LR', size='8,3')
                     tbar.set_svg(graph._repr_svg_())
-                tbar.set_description('avg reward: {:.2f}'.format(self.baseline))
+                if self.baseline:
+                    tbar.set_description('avg reward: {:.2f}'.format(self.baseline))
                 idx += 1
             self.validation()
             self.save()
-            tq.set_description('epoch {}, val_acc: {:.2f}, avg reward: {:.2f}' \
-                        .format(epoch, self.val_acc, self.baseline))
+            msg = 'epoch {}, val_acc: {:.2f}'.format(epoch, self.val_acc)
+            if self.baseline:
+                msg += ', avg reward: {:.2f}'.format(self.baseline)
+            tq.set_description(msg)
 
     def validation(self):
         if hasattr(self.val_data, 'reset'): self.val_data.reset()

{"golden_diff": "diff --git a/autogluon/contrib/enas/enas_scheduler.py b/autogluon/contrib/enas/enas_scheduler.py\n--- a/autogluon/contrib/enas/enas_scheduler.py\n+++ b/autogluon/contrib/enas/enas_scheduler.py\n@@ -126,12 +126,15 @@\n                     graph = self.supernet.graph\n                     graph.attr(rankdir='LR', size='8,3')\n                     tbar.set_svg(graph._repr_svg_())\n-                tbar.set_description('avg reward: {:.2f}'.format(self.baseline))\n+                if self.baseline:\n+                    tbar.set_description('avg reward: {:.2f}'.format(self.baseline))\n                 idx += 1\n             self.validation()\n             self.save()\n-            tq.set_description('epoch {}, val_acc: {:.2f}, avg reward: {:.2f}' \\\n-                        .format(epoch, self.val_acc, self.baseline))\n+            msg = 'epoch {}, val_acc: {:.2f}'.format(epoch, self.val_acc)\n+            if self.baseline:\n+                msg += ', avg reward: {:.2f}'.format(self.baseline)\n+            tq.set_description(msg)\n \n     def validation(self):\n         if hasattr(self.val_data, 'reset'): self.val_data.reset()\n", "issue": "About ENAS_Scheduler\nWhile using ENAS_Scheduler\r\nMaybe\uff1a\r\n```\r\ntbar.set_description('avg reward: {}'.format(self.baseline))\r\n```\r\nis better than \r\n\r\n```\r\ntbar.set_description('avg reward: {:.2f}'.format(self.baseline))\r\n```\r\nin run in enas_scheduler.py, because when self.baseline=None, it will report error\n", "before_files": [{"content": "import os\nimport pickle\nimport logging\nfrom collections import OrderedDict\nfrom multiprocessing.pool import ThreadPool\n\nimport mxnet as mx\n\nfrom ...searcher import RLSearcher\nfrom ...scheduler.resource import get_gpu_count, get_cpu_count\nfrom ...task.image_classification.dataset import get_built_in_dataset\nfrom ...task.image_classification.utils import *\nfrom ...utils import (mkdir, save, load, update_params, collect_params, DataLoader, tqdm, in_ipynb)\nfrom .enas_utils import *\n\n__all__ = ['ENAS_Scheduler']\n\nlogger = logging.getLogger(__name__)\n\nIMAGENET_TRAINING_SAMPLES = 1281167\n\nclass ENAS_Scheduler(object):\n    \"\"\"ENAS Scheduler, which automatically creates LSTM controller based on the search spaces.\n    \"\"\"\n    def __init__(self, supernet, train_set='imagenet', val_set=None,\n                 train_fn=default_train_fn, eval_fn=default_val_fn,\n                 train_args={}, val_args={}, reward_fn= default_reward_fn,\n                 num_gpus=0, num_cpus=4,\n                 batch_size=256, epochs=120, warmup_epochs=5,\n                 controller_lr=1e-3, controller_type='lstm',\n                 controller_batch_size=10, ema_baseline_decay=0.95,\n                 update_arch_frequency=20, checkname='./enas/checkpoint.ag',\n                 plot_frequency=0, **kwargs):\n        num_cpus = get_cpu_count() if num_cpus > get_cpu_count() else num_cpus\n        num_gpus = get_gpu_count() if num_gpus > get_gpu_count() else num_gpus\n        self.supernet = supernet\n        self.train_fn = train_fn\n        self.eval_fn = eval_fn\n        self.reward_fn = reward_fn\n        self.checkname = checkname\n        self.plot_frequency = plot_frequency\n        self.epochs = epochs\n        self.warmup_epochs = warmup_epochs\n        self.controller_batch_size = controller_batch_size\n        kwspaces = self.supernet.kwspaces\n\n        self.initialize_miscs(train_set, val_set, batch_size, num_cpus, num_gpus,\n                              train_args, val_args)\n\n        # create RL searcher/controller\n        self.baseline = None\n        self.ema_decay = ema_baseline_decay\n        self.searcher = RLSearcher(kwspaces, controller_type=controller_type,\n                                   prefetch=4, num_workers=4)\n        # controller setup\n        self.controller = self.searcher.controller\n        self.controller_optimizer = mx.gluon.Trainer(\n                self.controller.collect_params(), 'adam',\n                optimizer_params={'learning_rate': controller_lr})\n        self.update_arch_frequency = update_arch_frequency\n        self.val_acc = 0\n        # async controller sample\n        self._worker_pool = ThreadPool(2)\n        self._data_buffer = {}\n        self._rcvd_idx = 0\n        self._sent_idx = 0\n        self._timeout = 20\n        # logging history\n        self.training_history = []\n        self._prefetch_controller()\n\n    def initialize_miscs(self, train_set, val_set, batch_size, num_cpus, num_gpus,\n                         train_args, val_args):\n        \"\"\"Initialize framework related miscs, such as train/val data and train/val\n        function arguments.\n        \"\"\"\n        ctx = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu(0)]\n        self.supernet.collect_params().reset_ctx(ctx)\n        self.supernet.hybridize()\n        dataset_name = train_set\n\n        if isinstance(train_set, str):\n            train_set = get_built_in_dataset(dataset_name, train=True, batch_size=batch_size,\n                                             num_workers=num_cpus, shuffle=True)\n            val_set = get_built_in_dataset(dataset_name, train=False, batch_size=batch_size,\n                                           num_workers=num_cpus, shuffle=True)\n        if isinstance(train_set, gluon.data.Dataset):\n            self.train_data = DataLoader(\n                    train_set, batch_size=batch_size, shuffle=True,\n                    last_batch=\"discard\", num_workers=num_cpus)\n            # very important, make shuffle for training contoller\n            self.val_data = DataLoader(\n                    val_set, batch_size=batch_size, shuffle=True,\n                    num_workers=num_cpus, prefetch=0, sample_times=self.controller_batch_size)\n        else:\n            self.train_data = train_set\n            self.val_data = val_set\n        iters_per_epoch = len(self.train_data) if hasattr(self.train_data, '__len__') else \\\n                IMAGENET_TRAINING_SAMPLES // batch_size\n        self.train_args = init_default_train_args(batch_size, self.supernet, self.epochs, iters_per_epoch) \\\n                if len(train_args) == 0 else train_args\n        self.val_args = val_args\n        self.val_args['ctx'] = ctx\n        self.val_args['batch_fn'] = imagenet_batch_fn if dataset_name == 'imagenet' else default_batch_fn\n        self.train_args['ctx'] = ctx\n        self.train_args['batch_fn'] = imagenet_batch_fn if dataset_name == 'imagenet' else default_batch_fn\n        self.ctx = ctx\n\n    def run(self):\n        tq = tqdm(range(self.epochs))\n        for epoch in tq:\n            # for recordio data\n            if hasattr(self.train_data, 'reset'): self.train_data.reset()\n            tbar = tqdm(self.train_data)\n            idx = 0\n            for batch in tbar:\n                # sample network configuration\n                config = self.controller.pre_sample()[0]\n                self.supernet.sample(**config)\n                # self.train_fn(self.supernet, batch, **self.train_args)\n                self.train_fn(epoch, self.epochs, self.supernet, batch, **self.train_args)\n                mx.nd.waitall()\n                if epoch >= self.warmup_epochs and (idx % self.update_arch_frequency) == 0:\n                    self.train_controller()\n                if self.plot_frequency > 0 and idx % self.plot_frequency == 0 and in_ipynb():\n                    graph = self.supernet.graph\n                    graph.attr(rankdir='LR', size='8,3')\n                    tbar.set_svg(graph._repr_svg_())\n                tbar.set_description('avg reward: {:.2f}'.format(self.baseline))\n                idx += 1\n            self.validation()\n            self.save()\n            tq.set_description('epoch {}, val_acc: {:.2f}, avg reward: {:.2f}' \\\n                        .format(epoch, self.val_acc, self.baseline))\n\n    def validation(self):\n        if hasattr(self.val_data, 'reset'): self.val_data.reset()\n        # data iter, avoid memory leak\n        it = iter(self.val_data)\n        if hasattr(it, 'reset_sample_times'): it.reset_sample_times()\n        tbar = tqdm(it)\n        # update network arc\n        config = self.controller.inference()\n        self.supernet.sample(**config)\n        metric = mx.metric.Accuracy()\n        for batch in tbar:\n            self.eval_fn(self.supernet, batch, metric=metric, **self.val_args)\n            reward = metric.get()[1]\n            tbar.set_description('Val Acc: {}'.format(reward))\n\n        self.val_acc = reward\n        self.training_history.append(reward)\n\n    def _sample_controller(self):\n        assert self._rcvd_idx < self._sent_idx, \"rcvd_idx must be smaller than sent_idx\"\n        try:\n            ret = self._data_buffer.pop(self._rcvd_idx)\n            self._rcvd_idx += 1\n            return  ret.get(timeout=self._timeout)\n        except Exception:\n            self._worker_pool.terminate()\n            raise\n\n    def _prefetch_controller(self):\n        async_ret = self._worker_pool.apply_async(self._async_sample, ())\n        self._data_buffer[self._sent_idx] = async_ret\n        self._sent_idx += 1\n\n    def _async_sample(self):\n        with mx.autograd.record():\n            # sample controller_batch_size number of configurations\n            configs, log_probs, entropies = self.controller.sample(batch_size=self.controller_batch_size,\n                                                                   with_details=True)\n        return configs, log_probs, entropies\n\n    def train_controller(self):\n        \"\"\"Run multiple number of trials\n        \"\"\"\n        decay = self.ema_decay\n        if hasattr(self.val_data, 'reset'): self.val_data.reset()\n        # update \n        metric = mx.metric.Accuracy()\n        with mx.autograd.record():\n            # sample controller_batch_size number of configurations\n            configs, log_probs, entropies = self._sample_controller()\n            for i, batch in enumerate(self.val_data):\n                if i >= self.controller_batch_size: break\n                self.supernet.sample(**configs[i])\n                # schedule the training tasks and gather the reward\n                metric.reset()\n                self.eval_fn(self.supernet, batch, metric=metric, **self.val_args)\n                reward = metric.get()[1]\n                reward = self.reward_fn(reward, self.supernet)\n                self.baseline = reward if not self.baseline else self.baseline\n                # substract baseline\n                avg_rewards = mx.nd.array([reward - self.baseline],\n                                          ctx=self.controller.context)\n                # EMA baseline\n                self.baseline = decay * self.baseline + (1 - decay) * reward\n                # negative policy gradient\n                log_prob = log_probs[i]\n                log_prob = log_prob.sum()\n                loss = - log_prob * avg_rewards\n                loss = loss.sum()\n\n        # update\n        loss.backward()\n        self.controller_optimizer.step(self.controller_batch_size)\n        self._prefetch_controller()\n\n    def load(self, checkname=None):\n        checkname = checkname if checkname else self.checkname\n        state_dict = load(checkname)\n        self.load_state_dict(state_dict)\n\n    def save(self, checkname=None):\n        checkname = checkname if checkname else self.checkname\n        mkdir(os.path.dirname(checkname))\n        save(self.state_dict(), checkname)\n\n    def state_dict(self, destination=None):\n        if destination is None:\n            destination = OrderedDict()\n            destination._metadata = OrderedDict()\n        destination['supernet_params'] = collect_params(self.supernet)\n        destination['controller_params'] = collect_params(self.controller)\n        destination['training_history'] = self.training_history\n        return destination\n\n    def load_state_dict(self, state_dict):\n        update_params(self.supernet, state_dict['supernet_params'], ctx=self.ctx)\n        update_params(self.controller, state_dict['controller_params'], ctx=self.controller.context)\n        self.training_history = state_dict['training_history']\n", "path": "autogluon/contrib/enas/enas_scheduler.py"}], "after_files": [{"content": "import os\nimport pickle\nimport logging\nfrom collections import OrderedDict\nfrom multiprocessing.pool import ThreadPool\n\nimport mxnet as mx\n\nfrom ...searcher import RLSearcher\nfrom ...scheduler.resource import get_gpu_count, get_cpu_count\nfrom ...task.image_classification.dataset import get_built_in_dataset\nfrom ...task.image_classification.utils import *\nfrom ...utils import (mkdir, save, load, update_params, collect_params, DataLoader, tqdm, in_ipynb)\nfrom .enas_utils import *\n\n__all__ = ['ENAS_Scheduler']\n\nlogger = logging.getLogger(__name__)\n\nIMAGENET_TRAINING_SAMPLES = 1281167\n\nclass ENAS_Scheduler(object):\n    \"\"\"ENAS Scheduler, which automatically creates LSTM controller based on the search spaces.\n    \"\"\"\n    def __init__(self, supernet, train_set='imagenet', val_set=None,\n                 train_fn=default_train_fn, eval_fn=default_val_fn,\n                 train_args={}, val_args={}, reward_fn= default_reward_fn,\n                 num_gpus=0, num_cpus=4,\n                 batch_size=256, epochs=120, warmup_epochs=5,\n                 controller_lr=1e-3, controller_type='lstm',\n                 controller_batch_size=10, ema_baseline_decay=0.95,\n                 update_arch_frequency=20, checkname='./enas/checkpoint.ag',\n                 plot_frequency=0, **kwargs):\n        num_cpus = get_cpu_count() if num_cpus > get_cpu_count() else num_cpus\n        num_gpus = get_gpu_count() if num_gpus > get_gpu_count() else num_gpus\n        self.supernet = supernet\n        self.train_fn = train_fn\n        self.eval_fn = eval_fn\n        self.reward_fn = reward_fn\n        self.checkname = checkname\n        self.plot_frequency = plot_frequency\n        self.epochs = epochs\n        self.warmup_epochs = warmup_epochs\n        self.controller_batch_size = controller_batch_size\n        kwspaces = self.supernet.kwspaces\n\n        self.initialize_miscs(train_set, val_set, batch_size, num_cpus, num_gpus,\n                              train_args, val_args)\n\n        # create RL searcher/controller\n        self.baseline = None\n        self.ema_decay = ema_baseline_decay\n        self.searcher = RLSearcher(kwspaces, controller_type=controller_type,\n                                   prefetch=4, num_workers=4)\n        # controller setup\n        self.controller = self.searcher.controller\n        self.controller_optimizer = mx.gluon.Trainer(\n                self.controller.collect_params(), 'adam',\n                optimizer_params={'learning_rate': controller_lr})\n        self.update_arch_frequency = update_arch_frequency\n        self.val_acc = 0\n        # async controller sample\n        self._worker_pool = ThreadPool(2)\n        self._data_buffer = {}\n        self._rcvd_idx = 0\n        self._sent_idx = 0\n        self._timeout = 20\n        # logging history\n        self.training_history = []\n        self._prefetch_controller()\n\n    def initialize_miscs(self, train_set, val_set, batch_size, num_cpus, num_gpus,\n                         train_args, val_args):\n        \"\"\"Initialize framework related miscs, such as train/val data and train/val\n        function arguments.\n        \"\"\"\n        ctx = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu(0)]\n        self.supernet.collect_params().reset_ctx(ctx)\n        self.supernet.hybridize()\n        dataset_name = train_set\n\n        if isinstance(train_set, str):\n            train_set = get_built_in_dataset(dataset_name, train=True, batch_size=batch_size,\n                                             num_workers=num_cpus, shuffle=True)\n            val_set = get_built_in_dataset(dataset_name, train=False, batch_size=batch_size,\n                                           num_workers=num_cpus, shuffle=True)\n        if isinstance(train_set, gluon.data.Dataset):\n            self.train_data = DataLoader(\n                    train_set, batch_size=batch_size, shuffle=True,\n                    last_batch=\"discard\", num_workers=num_cpus)\n            # very important, make shuffle for training contoller\n            self.val_data = DataLoader(\n                    val_set, batch_size=batch_size, shuffle=True,\n                    num_workers=num_cpus, prefetch=0, sample_times=self.controller_batch_size)\n        else:\n            self.train_data = train_set\n            self.val_data = val_set\n        iters_per_epoch = len(self.train_data) if hasattr(self.train_data, '__len__') else \\\n                IMAGENET_TRAINING_SAMPLES // batch_size\n        self.train_args = init_default_train_args(batch_size, self.supernet, self.epochs, iters_per_epoch) \\\n                if len(train_args) == 0 else train_args\n        self.val_args = val_args\n        self.val_args['ctx'] = ctx\n        self.val_args['batch_fn'] = imagenet_batch_fn if dataset_name == 'imagenet' else default_batch_fn\n        self.train_args['ctx'] = ctx\n        self.train_args['batch_fn'] = imagenet_batch_fn if dataset_name == 'imagenet' else default_batch_fn\n        self.ctx = ctx\n\n    def run(self):\n        tq = tqdm(range(self.epochs))\n        for epoch in tq:\n            # for recordio data\n            if hasattr(self.train_data, 'reset'): self.train_data.reset()\n            tbar = tqdm(self.train_data)\n            idx = 0\n            for batch in tbar:\n                # sample network configuration\n                config = self.controller.pre_sample()[0]\n                self.supernet.sample(**config)\n                # self.train_fn(self.supernet, batch, **self.train_args)\n                self.train_fn(epoch, self.epochs, self.supernet, batch, **self.train_args)\n                mx.nd.waitall()\n                if epoch >= self.warmup_epochs and (idx % self.update_arch_frequency) == 0:\n                    self.train_controller()\n                if self.plot_frequency > 0 and idx % self.plot_frequency == 0 and in_ipynb():\n                    graph = self.supernet.graph\n                    graph.attr(rankdir='LR', size='8,3')\n                    tbar.set_svg(graph._repr_svg_())\n                if self.baseline:\n                    tbar.set_description('avg reward: {:.2f}'.format(self.baseline))\n                idx += 1\n            self.validation()\n            self.save()\n            msg = 'epoch {}, val_acc: {:.2f}'.format(epoch, self.val_acc)\n            if self.baseline:\n                msg += ', avg reward: {:.2f}'.format(self.baseline)\n            tq.set_description(msg)\n\n    def validation(self):\n        if hasattr(self.val_data, 'reset'): self.val_data.reset()\n        # data iter, avoid memory leak\n        it = iter(self.val_data)\n        if hasattr(it, 'reset_sample_times'): it.reset_sample_times()\n        tbar = tqdm(it)\n        # update network arc\n        config = self.controller.inference()\n        self.supernet.sample(**config)\n        metric = mx.metric.Accuracy()\n        for batch in tbar:\n            self.eval_fn(self.supernet, batch, metric=metric, **self.val_args)\n            reward = metric.get()[1]\n            tbar.set_description('Val Acc: {}'.format(reward))\n\n        self.val_acc = reward\n        self.training_history.append(reward)\n\n    def _sample_controller(self):\n        assert self._rcvd_idx < self._sent_idx, \"rcvd_idx must be smaller than sent_idx\"\n        try:\n            ret = self._data_buffer.pop(self._rcvd_idx)\n            self._rcvd_idx += 1\n            return  ret.get(timeout=self._timeout)\n        except Exception:\n            self._worker_pool.terminate()\n            raise\n\n    def _prefetch_controller(self):\n        async_ret = self._worker_pool.apply_async(self._async_sample, ())\n        self._data_buffer[self._sent_idx] = async_ret\n        self._sent_idx += 1\n\n    def _async_sample(self):\n        with mx.autograd.record():\n            # sample controller_batch_size number of configurations\n            configs, log_probs, entropies = self.controller.sample(batch_size=self.controller_batch_size,\n                                                                   with_details=True)\n        return configs, log_probs, entropies\n\n    def train_controller(self):\n        \"\"\"Run multiple number of trials\n        \"\"\"\n        decay = self.ema_decay\n        if hasattr(self.val_data, 'reset'): self.val_data.reset()\n        # update \n        metric = mx.metric.Accuracy()\n        with mx.autograd.record():\n            # sample controller_batch_size number of configurations\n            configs, log_probs, entropies = self._sample_controller()\n            for i, batch in enumerate(self.val_data):\n                if i >= self.controller_batch_size: break\n                self.supernet.sample(**configs[i])\n                # schedule the training tasks and gather the reward\n                metric.reset()\n                self.eval_fn(self.supernet, batch, metric=metric, **self.val_args)\n                reward = metric.get()[1]\n                reward = self.reward_fn(reward, self.supernet)\n                self.baseline = reward if not self.baseline else self.baseline\n                # substract baseline\n                avg_rewards = mx.nd.array([reward - self.baseline],\n                                          ctx=self.controller.context)\n                # EMA baseline\n                self.baseline = decay * self.baseline + (1 - decay) * reward\n                # negative policy gradient\n                log_prob = log_probs[i]\n                log_prob = log_prob.sum()\n                loss = - log_prob * avg_rewards\n                loss = loss.sum()\n\n        # update\n        loss.backward()\n        self.controller_optimizer.step(self.controller_batch_size)\n        self._prefetch_controller()\n\n    def load(self, checkname=None):\n        checkname = checkname if checkname else self.checkname\n        state_dict = load(checkname)\n        self.load_state_dict(state_dict)\n\n    def save(self, checkname=None):\n        checkname = checkname if checkname else self.checkname\n        mkdir(os.path.dirname(checkname))\n        save(self.state_dict(), checkname)\n\n    def state_dict(self, destination=None):\n        if destination is None:\n            destination = OrderedDict()\n            destination._metadata = OrderedDict()\n        destination['supernet_params'] = collect_params(self.supernet)\n        destination['controller_params'] = collect_params(self.controller)\n        destination['training_history'] = self.training_history\n        return destination\n\n    def load_state_dict(self, state_dict):\n        update_params(self.supernet, state_dict['supernet_params'], ctx=self.ctx)\n        update_params(self.controller, state_dict['controller_params'], ctx=self.controller.context)\n        self.training_history = state_dict['training_history']\n", "path": "autogluon/contrib/enas/enas_scheduler.py"}]}