Function load_and_assign_ckpt

tensorlayer/files/utils.py:2823–2859

Load weights by name from a given file of ckpt format Parameters ---------- model_dir : str Filename to which the weights will be loaded, should be of ckpt format. Examples: model_dir = /root/cnn_model/ network : Model TL model. skip : bool If 'sk

(model_dir, network=None, skip=True)

def load_and_assign_ckpt(model_dir, network=None, skip=True):
    """Load weights by name from a given file of ckpt format

    Parameters
    ----------
    model_dir : str
        Filename to which the weights will be loaded, should be of ckpt format.
        Examples: model_dir = /root/cnn_model/
    network : Model
        TL model.
    skip : bool
        If 'skip' == True, loaded weights whose name is not found in 'weights' will be skipped. If 'skip' is False,
        error will be raised when mismatch is found. Default False.

    Returns
    -------

    """
    model_path, filename = check_ckpt_file(model_dir)

    reader = pywrap_tensorflow.NewCheckpointReader(model_path)
    var_to_shape_map = reader.get_variable_to_shape_map()

    net_weights_name = [w.name for w in network.all_weights]

    for key in var_to_shape_map:
        if key not in net_weights_name:
            if skip:
                logging.warning("Weights named '%s' not found in network. Skip it." % key)
            else:
                raise RuntimeError(
                    "Weights named '%s' not found in network. Hint: set argument skip=Ture "
                    "if you want to skip redundant or mismatch weights." % key
                )
        else:
            assign_tf_variable(network.all_weights[net_weights_name.index(key)], reader.get_tensor(key))
    logging.info("[*] Model restored from ckpt %s" % filename)


def ckpt_to_npz_dict(model_dir, save_name='model.npz', rename_key=False):