Class TrainTowerContext

tensorpack/tfutils/tower.py:163–199

class TrainTowerContext(BaseTowerContext):

    def __init__(self, ns_name, vs_name='', index=0, total=1):
        """
        Args:
            index (int): index of this tower, only used in training.
            total (int): total number of towers to be built.
        """
        super(TrainTowerContext, self).__init__(ns_name, vs_name)
        self._is_training = True

        self.index = int(index)
        self.total = int(total)
        if self.index > 0:
            assert self.total > self.index, "(index, total) = ({}, {})".format(self.index, self.total)

        vs = tf.get_variable_scope()
        assert vs.name == '', "Cannot nest TrainTowerContext with an existing variable scope!"
        if vs_name:
            assert not vs.reuse, \
                "Cannot create tower {} with vs_name={} under reuse=True!".format(ns_name, vs_name)
        self._original_vs_reuse = vs.reuse

    @property
    def is_main_training_tower(self):
        return self.index == 0

    @property
    def has_own_variables(self):
        if self._original_vs_reuse:
            return False
        return self.index == 0 or len(self._vs_name) > 0

    def _keys_to_freeze(self):
        if self.index == 0:
            return []
        return [tf.GraphKeys.SUMMARIES, MOVING_SUMMARY_OPS_KEY]


class PredictTowerContext(BaseTowerContext):