Args: config (PredictConfig): the config to use. towers: a list of relative GPU id.
(self, config, towers)
| 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 |
no test coverage detected