Function freeze_variables

tensorpack/tfutils/varreplace.py:61–129

Return a context to freeze variables, by wrapping ``tf.get_variable`` with a custom getter. It works by either applying ``tf.stop_gradient`` on the variables, or keeping them out of the ``TRAINABLE_VARIABLES`` collection, or both. Both options have their own pros and cons.

(stop_gradient=True, skip_collection=False)

def freeze_variables(stop_gradient=True, skip_collection=False):
    """
    Return a context to freeze variables,
    by wrapping ``tf.get_variable`` with a custom getter.
    It works by either applying ``tf.stop_gradient`` on the variables,
    or keeping them out of the ``TRAINABLE_VARIABLES`` collection, or
    both. Both options have their own pros and cons.

    Example:
        .. code-block:: python

            from tensorpack.tfutils import varreplace
            with varreplace.freeze_variable(stop_gradient=False, skip_collection=True):
                x = FullyConnected('fc', x, 1000)   # fc/* will not be trained

    Args:
        stop_gradient (bool): if True, variables returned from `get_variable`
            will be wrapped with `tf.stop_gradient`.

            Note that the created variables may still have gradient when accessed
            by other approaches (e.g. by name, or by collection).
            For example, they may still have a gradient in weight decay.
            Also note that this makes `tf.get_variable` returns a Tensor instead of a Variable,
            which may break existing contract.
            Therefore, it's recommended to use the `skip_collection` option instead.
        skip_collection (bool): if True, do not add the variable to
            ``TRAINABLE_VARIABLES`` collection, but to ``MODEL_VARIABLES``
            collection. As a result they will not be trained by default.

    Note:

    `stop_gradient` only stops variables returned by `get_variable` **within the context** to
    contribute no gradient in this context. Therefore it may not completely freeze the variables.
    For example:

        1. If a variable is created, or reused outside of the context, it can still contribute to the
           gradient of other tensors.
        2. If a freezed variable is accessed by other approaches (e.g., by names, by collections),
           it can still contribute to the gradient of other tensors.
           For example, weight decay cannot be stopped by a `stop_gradient` context.

    `skip_collection` has to be used the first time the variable is created.
    Once `skip_collection` is used, the variable is not a trainable variable anymore,
    and will be completely freezed from gradient update in tensorpack's single-cost trainer.

    Choose the option carefully depend on what you need.
    """
    def custom_getter(getter, *args, **kwargs):
        trainable = kwargs.get('trainable', True)
        name = args[0] if len(args) else kwargs.get('name')
        if skip_collection:
            kwargs['trainable'] = False
        v = getter(*args, **kwargs)
        # do not perform unnecessary changes if it's not originally trainable
        # otherwise the variable may get added to MODEL_VARIABLES twice
        if trainable and skip_collection:
            if isinstance(v, tf.Variable):
                tf.add_to_collection(tf.GraphKeys.MODEL_VARIABLES, v)