Method setup_graph

tensorpack/train/tower.py:197–220

Responsible for building the main training graph for single-cost training. Args: input_signature ([TensorSpec]): list of TensorSpec that describe the inputs input (InputSource): an InputSource which has to match the input signature get_cost_fn ([

(self, input_signature, input, get_cost_fn, get_opt_fn)

    @call_only_once
    def setup_graph(self, input_signature, input, get_cost_fn, get_opt_fn):
        """
        Responsible for building the main training graph for single-cost training.

        Args:
            input_signature ([TensorSpec]): list of TensorSpec that describe the inputs
            input (InputSource): an InputSource which has to match the input signature
            get_cost_fn ([tf.Tensor] -> tf.Tensor): callable, takes some input tensors and return a cost tensor.
            get_opt_fn (-> tf.train.Optimizer): callable which returns an
                optimizer. Will only be called once.

        Note:
            `get_cost_fn` will be part of the tower function.
            It must follows the `rules of tower function.
            <http://tensorpack.readthedocs.io/tutorial/trainer.html#tower-trainer>`_.
        """
        get_cost_fn = TowerFunc(get_cost_fn, input_signature)
        get_opt_fn = memoized(get_opt_fn)
        self.tower_func = get_cost_fn

        # TODO setup may want to register monitor as well??
        input_callbacks = self._setup_input(input_signature, input)
        train_callbacks = self._setup_graph(input, get_cost_fn, get_opt_fn)
        self.register_callback(input_callbacks + train_callbacks)

    @abstractmethod
    def _setup_graph(self, input, get_cost_fn, get_opt_fn):