(self)
| 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) |
nothing calls this directly
no test coverage detected