Class DataParallelInferenceRunner

tensorpack/callbacks/inference_runner.py:183–307

Inference with data-parallel support on multiple GPUs. It will build one predict tower on each GPU, and run prediction with a large total batch in parallel on all GPUs. It will run the remainder (when the total size of input is not a multiple of #GPU) sequentially.

class DataParallelInferenceRunner(InferenceRunnerBase):
    """
    Inference with data-parallel support on multiple GPUs.
    It will build one predict tower on each GPU, and run prediction
    with a large total batch in parallel on all GPUs.
    It will run the remainder (when the total size of input is not a multiple of #GPU)
    sequentially.
    """
    def __init__(self, input, infs, gpus, tower_name='InferenceTower', tower_func=None):
        """
        Args:
            input (DataFlow or QueueInput)
            gpus (int or list[int]): #gpus, or list of GPU id
            tower_name (str): the name scope of the tower to build.
                If multiple InferenceRunner are used, each needs a different tower_name.
            tower_func (tfutils.TowerFunc or None): the tower function to be used to build the graph.
                The tower function will be called under a `training=False` TowerContext.
                The default is `trainer.tower_func`,
                but you can change it to a different tower function
                if you need to inference with several different models.
        """
        if isinstance(gpus, int):
            gpus = list(range(gpus))
        self._devices = [_device_from_int(k) for k in gpus]
        self._tower_names = ['{}{}'.format(tower_name, k) for k in range(len(gpus))]

        if isinstance(input, DataFlow):
            input = QueueInput(input)
        assert isinstance(input, QueueInput), input
        super(DataParallelInferenceRunner, self).__init__(input, infs)
        assert self._size > 0, "Input for DataParallelInferenceRunner must have a size!"

        self._hooks = []
        self._hooks_parallel = []
        self._tower_func = tower_func

    def _setup_graph(self):
        self._handles = []
        if self._tower_func is None:
            assert self.trainer.tower_func is not None, "You must set tower_func of the trainer to use InferenceRunner!"
            self._tower_func = self.trainer.tower_func

        input_callbacks = self._input_source.setup(self._tower_func.input_signature)
        with tf.variable_scope(tf.get_variable_scope(), reuse=True):
            for idx, dev in enumerate(self._devices):
                vs_name = self.trainer._vs_name_for_predictor(idx)
                with tf.device(dev), PredictTowerContext(
                        self._tower_names[idx], vs_name=vs_name):
                    logger.info("[InferenceRunner] Building tower '{}' on device {} {}...".format(
                        self._tower_names[idx], dev,
                        "with variable scope '{}'".format(vs_name) if vs_name else ''))
                    # TODO log for tower creation, here or in tower.py?
                    self._tower_func(*self._input_source.get_input_tensors())
                    self._handles.append(self._tower_func.towers[-1])

        # setup callbacks and hooks
        self._input_callbacks = Callbacks(input_callbacks)