A multi-tower multi-GPU predictor. It builds one predictor for each tower.
| 14 | |
| 15 | |
| 16 | class MultiTowerOfflinePredictor(OnlinePredictor): |
| 17 | """ A multi-tower multi-GPU predictor. |
| 18 | It builds one predictor for each tower. |
| 19 | """ |
| 20 | |
| 21 | def __init__(self, config, towers): |
| 22 | """ |
| 23 | Args: |
| 24 | config (PredictConfig): the config to use. |
| 25 | towers: a list of relative GPU id. |
| 26 | """ |
| 27 | assert len(towers) > 0 |
| 28 | self.graph = config._maybe_create_graph() |
| 29 | self.predictors = [] |
| 30 | self.return_input = config.return_input |
| 31 | with self.graph.as_default(): |
| 32 | handles = [] |
| 33 | |
| 34 | input = PlaceholderInput() |
| 35 | input.setup(config.input_signature) |
| 36 | |
| 37 | for idx, t in enumerate(towers): |
| 38 | tower_name = 'tower' + str(t) |
| 39 | |
| 40 | device = '/gpu:{}'.format(t) |
| 41 | with tf.variable_scope(tf.get_variable_scope(), reuse=idx > 0), \ |
| 42 | tf.device(device), \ |
| 43 | PredictTowerContext(tower_name): |
| 44 | logger.info("Building graph for predict tower '{}' on device {} ...".format(tower_name, device)) |
| 45 | config.tower_func(*input.get_input_tensors()) |
| 46 | handles.append(config.tower_func.towers[-1]) |
| 47 | |
| 48 | config.session_init._setup_graph() |
| 49 | self.sess = config.session_creator.create_session() |
| 50 | config.session_init._run_init(self.sess) |
| 51 | |
| 52 | for h in handles: |
| 53 | input_tensors = h.get_tensors(config.input_names) |
| 54 | output_tensors = h.get_tensors(config.output_names) |
| 55 | self.predictors.append(OnlinePredictor( |
| 56 | input_tensors, output_tensors, config.return_input, self.sess)) |
| 57 | |
| 58 | def _do_call(self, dp): |
| 59 | # use the first tower for compatible PredictorBase interface |
| 60 | return self.predictors[0]._do_call(dp) |
| 61 | |
| 62 | def get_predictor(self, n): |
| 63 | """ |
| 64 | Returns: |
| 65 | OnlinePredictor: the nth predictor on the nth tower. |
| 66 | """ |
| 67 | l = len(self.predictors) |
| 68 | if n >= l: |
| 69 | logger.warn("n > #towers, will assign predictor to GPU by round-robin") |
| 70 | return [self.predictors[k % l] for k in range(n)] |
| 71 | |
| 72 | def get_predictors(self): |
| 73 | """ |