Class DictRestore

tensorpack/tfutils/sessinit.py:196–241

Restore variables from a dictionary.

class DictRestore(SessionInit):
    """
    Restore variables from a dictionary.
    """

    def __init__(self, variable_dict, ignore_mismatch=False):
        """
        Args:
            variable_dict (dict): a dict of {name: value}
            ignore_mismatch (bool): ignore failures when the value and the
                variable does not match in their shapes.
                If False, it will throw exception on such errors.
                If True, it will only print a warning.
        """
        assert isinstance(variable_dict, dict), type(variable_dict)
        # use varname (with :0) for consistency
        self._prms = {get_op_tensor_name(n)[1]: v for n, v in six.iteritems(variable_dict)}
        self._ignore_mismatch = ignore_mismatch

    def _run_init(self, sess):
        variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
        variable_names_list = [k.name for k in variables]

        variable_names = set(variable_names_list)
        param_names = set(six.iterkeys(self._prms))

        # intersect has the original ordering of variables
        intersect = [v for v in variable_names_list if v in param_names]

        # use opname (without :0) for clarity in logging
        logger.info("Variables to restore from dict: {}".format(
            ', '.join(get_op_tensor_name(x)[0] for x in intersect)))

        mismatch = MismatchLogger('graph', 'dict')
        for k in sorted(variable_names - param_names):
            if not is_training_name(k):
                mismatch.add(k)
        mismatch.log()
        mismatch = MismatchLogger('dict', 'graph')
        for k in sorted(param_names - variable_names):
            mismatch.add(k)
        mismatch.log()

        upd = SessionUpdate(sess, [v for v in variables if v.name in intersect], ignore_mismatch=self._ignore_mismatch)
        logger.info("Restoring {} variables from dict ...".format(len(intersect)))
        upd.update({name: value for name, value in six.iteritems(self._prms) if name in intersect})


class ChainInit(SessionInit):