Method __init__

tensorpack/predict/multigpu.py:90–121

Args: config (PredictConfig): the config to use. towers: a list of relative GPU id.

(self, config, towers)

    """

    def __init__(self, config, towers):
        """
        Args:
            config (PredictConfig): the config to use.
            towers: a list of relative GPU id.
        """
        self.graph = config._maybe_create_graph()
        with self.graph.as_default():
            input_tensors = []
            output_tensors = []

            for idx, t in enumerate(towers):
                tower_name = 'tower' + str(t)

                new_sig = [tf.TensorSpec(dtype=p.dtype, shape=p.shape, name=tower_name + '_' + p.name)
                           for p in config.input_signature]
                input = PlaceholderInput()
                input.setup(new_sig)

                with tf.variable_scope(tf.get_variable_scope(), reuse=idx > 0), \
                        tf.device('/gpu:{}'.format(t)), \
                        PredictTowerContext(tower_name):
                    config.tower_func(*input.get_input_tensors())
                    h = config.tower_func.towers[-1]
                    input_tensors.extend(h.get_tensors(config.input_names))
                    output_tensors.extend(h.get_tensors(config.output_names))

            config.session_init._setup_graph()
            sess = config.session_creator.create_session()
            config.session_init._run_init(sess)
            super(DataParallelOfflinePredictor, self).__init__(
                input_tensors, output_tensors, config.return_input, sess)