Method get_predictor

tensorpack/train/tower.py:83–142

This method will build the trainer's tower function under ``TowerContext(is_training=False)``, and returns a callable predictor with input placeholders & output tensors in this tower. This method handles the common case where you inference with the same tower function

(self, input_names, output_names, device=0)

        return self.tower_func.towers

    def get_predictor(self, input_names, output_names, device=0):
        """
        This method will build the trainer's tower function under ``TowerContext(is_training=False)``,
        and returns a callable predictor with input placeholders & output tensors in this tower.

        This method handles the common case where you inference with the same tower function
        you provide to the trainer.
        If you want to do inference with a different tower function, you can always build the tower by yourself,
        under a "reuse" variable scope and a `TowerContext(is_training=False)`.

        Args:
            input_names (list): list of input names, matching the inputs declared for the trainer.
            output_names(list): list of tensor names without the tower prefix.
            device (int): build the predictor on device '/gpu:{device}' or use -1 for '/cpu:0'.

        Returns:
            an :class:`OnlinePredictor`.

        Example:

        .. code-block:: none

            # in the graph:
            interesting_tensor = tf.identity(x, name='fun')
            # in _setup_graph callback method:
            self._predictor = self.trainer.get_predictor(['input1', 'input2'], ['fun'])
            # After session is initialized (see Tutorials - Write a Callback), can use it by:
            outputs = self._predictor(input1, input2)

        The CycleGAN example and DQN example have more concrete use of this method.
        """
        assert self.tower_func is not None, "Must set tower_func on the trainer to use get_predictor()!"
        tower_name = 'tower-pred-{}'.format(device) if device >= 0 else 'tower-pred-cpu'
        device_id = device
        device = '/gpu:{}'.format(device_id) if device_id >= 0 else '/cpu:0'

        try:
            tower = self.tower_func.towers[tower_name]
            assert tower is not None, "This is a bug!"
        except KeyError:
            tower = None

        if tower is None:
            input = PlaceholderInput()
            input.setup(self.input_signature)

            vs_name = self._vs_name_for_predictor(device_id)
            with tfv1.variable_scope(tfv1.get_variable_scope(), reuse=True), \
                    tf.device(device), PredictTowerContext(
                        tower_name, vs_name=vs_name):
                logger.info("Building graph for predict tower '{}' on device {} {}...".format(
                    tower_name, device,
                    "with variable scope '{}'".format(vs_name) if vs_name else ''))
                self.tower_func(*input.get_input_tensors())
            tower = self.tower_func.towers[tower_name]
        input_tensors = tower.get_tensors(input_names)
        output_tensors = tower.get_tensors(output_names)
        predictor = OnlinePredictor(input_tensors, output_tensors)