Class TowerFunc

tensorpack/tfutils/tower.py:251–309

A tower function (see `tutorial on tower function `_) It keeps track of the name scope, variable scope and input/output tensors each time the function is called. :class:`TowerTrainer` needs this so

class TowerFunc(object):
    """
    A tower function (see
    `tutorial on tower function
    <http://tensorpack.readthedocs.io/tutorial/extend/trainer.html#tower-trainer>`_)
    It keeps track of the name scope, variable scope and input/output tensors
    each time the function is called.

    :class:`TowerTrainer` needs this so that it knows how to build a predictor.

    Conceptually, this class is roughly equivalent to `tf.function` with input signature, introduced in TF 2.0.
    """

    def __init__(self, tower_fn, input_signature):
        """
        Args:
            tower_func: a function which builds one tower in the graph.
                It takes several input tensors and could return anything.
            input_signature ([TensorSpec]): list of :class:`tf.TensorSpec`.
                They are used to figure out the names for the input tensors.
        """
        assert callable(tower_fn), tower_fn
        self._inputs_names = [k.name for k in input_signature]
        assert len(set(self._inputs_names)) == len(self._inputs_names), \
            "Duplicated names in input_signature! " + str(self._inputs_names)
        for name in self._inputs_names:
            if any(k in name for k in [':', '/', ' ']):
                raise ValueError("Invalid input name: '{}'".format(name))
        self._tower_fn = tower_fn
        self._input_signature = input_signature

        self._handles = []

    def __new__(cls, tower_fn, _):
        # to avoid double-wrapping a function
        if isinstance(tower_fn, TowerFunc):
            return tower_fn
        else:
            return super(TowerFunc, cls).__new__(cls)

    def __call__(self, *args):
        ctx = get_current_tower_context()
        assert ctx is not None, "Function must be called under TowerContext!"
        output = self._tower_fn(*args)
        handle = TowerTensorHandle(ctx, args, output, self._input_signature)
        self._handles.append(handle)
        return output

    @property
    def towers(self):
        """
        TowerTensorHandles: a :class:`TowerTensorHandles` object, that can
            access the tower handles by either indices or names.
        """
        return TowerTensorHandles(self._handles)

    @property
    def input_signature(self):