A callback that runs evaluation once a while. It supports multi-gpu evaluation.
| 209 | |
| 210 | |
| 211 | class EvalCallback(Callback): |
| 212 | """ |
| 213 | A callback that runs evaluation once a while. |
| 214 | It supports multi-gpu evaluation. |
| 215 | """ |
| 216 | |
| 217 | _chief_only = False |
| 218 | |
| 219 | def __init__(self, eval_dataset, in_names, out_names, output_dir): |
| 220 | self._eval_dataset = eval_dataset |
| 221 | self._in_names, self._out_names = in_names, out_names |
| 222 | self._output_dir = output_dir |
| 223 | |
| 224 | def _setup_graph(self): |
| 225 | num_gpu = cfg.TRAIN.NUM_GPUS |
| 226 | if cfg.TRAINER == 'replicated': |
| 227 | # TF bug in version 1.11, 1.12: https://github.com/tensorflow/tensorflow/issues/22750 |
| 228 | buggy_tf = get_tf_version_tuple() in [(1, 11), (1, 12)] |
| 229 | |
| 230 | # Use two predictor threads per GPU to get better throughput |
| 231 | self.num_predictor = num_gpu if buggy_tf else num_gpu * 2 |
| 232 | self.predictors = [self._build_predictor(k % num_gpu) for k in range(self.num_predictor)] |
| 233 | self.dataflows = [get_eval_dataflow(self._eval_dataset, |
| 234 | shard=k, num_shards=self.num_predictor) |
| 235 | for k in range(self.num_predictor)] |
| 236 | else: |
| 237 | # Only eval on the first machine, |
| 238 | # Because evaluation assumes that all horovod workers share the filesystem. |
| 239 | # Alternatively, can eval on all ranks and use allgather, but allgather sometimes hangs |
| 240 | self._horovod_run_eval = hvd.rank() == hvd.local_rank() |
| 241 | if self._horovod_run_eval: |
| 242 | self.predictor = self._build_predictor(0) |
| 243 | self.dataflow = get_eval_dataflow(self._eval_dataset, |
| 244 | shard=hvd.local_rank(), num_shards=hvd.local_size()) |
| 245 | |
| 246 | self.barrier = hvd.allreduce(tf.random_normal(shape=[1])) |
| 247 | |
| 248 | def _build_predictor(self, idx): |
| 249 | return self.trainer.get_predictor(self._in_names, self._out_names, device=idx) |
| 250 | |
| 251 | def _before_train(self): |
| 252 | eval_period = cfg.TRAIN.EVAL_PERIOD |
| 253 | self.epochs_to_eval = set() |
| 254 | for k in itertools.count(1): |
| 255 | if k * eval_period > self.trainer.max_epoch: |
| 256 | break |
| 257 | self.epochs_to_eval.add(k * eval_period) |
| 258 | self.epochs_to_eval.add(self.trainer.max_epoch) |
| 259 | logger.info("[EvalCallback] Will evaluate every {} epochs".format(eval_period)) |
| 260 | |
| 261 | def _eval(self): |
| 262 | logdir = self._output_dir |
| 263 | if cfg.TRAINER == 'replicated': |
| 264 | all_results = multithread_predict_dataflow(self.dataflows, self.predictors) |
| 265 | else: |
| 266 | filenames = [os.path.join( |
| 267 | logdir, 'outputs{}-part{}.json'.format(self.global_step, rank) |
| 268 | ) for rank in range(hvd.local_size())] |