Merge pull request #1075 from Laicheng0830/Add_load_ckpt

add load ckpt weights

Author: Laicheng0830

examples/pretrained_cnn/tutorial_load_ckpt_weights_to_tensorlayer.py

@@ -0,0 +1,70 @@
+#! /usr/bin/python
+# -*- coding: utf-8 -*-
+
+import tensorlayer as tl
+from tensorlayer.layers import (Input, Conv2d, Flatten, Dense, MaxPool2d)
+from tensorlayer.models import Model
+from tensorlayer.files import maybe_download_and_extract
+import numpy as np
+import tensorflow as tf
+
+filename = 'ckpt_parameters.zip'
+url_score = 'https://media.githubusercontent.com/media/tensorlayer/pretrained-models/master/models/'
+
+# download weights
+down_file = tl.files.maybe_download_and_extract(
+    filename=filename, working_directory='model/', url_source=url_score, extract=True
+)
+
+model_file = 'model/ckpt_parameters'
+
+# ckpt to npz, rename_key used to match TL naming rule
+tl.files.ckpt_to_npz_dict(model_file, rename_key=True)
+weights = np.load('model.npz', allow_pickle=True)
+
+# View the parameters and weights shape
+for key in weights.keys():
+    print(key, weights[key].shape)
+
+
+# build model
+def create_model(inputs_shape):
+    W_init = tl.initializers.truncated_normal(stddev=5e-2)
+    W_init2 = tl.initializers.truncated_normal(stddev=0.04)
+    ni = Input(inputs_shape)
+    nn = Conv2d(64, (3, 3), (1, 1), padding='SAME', act=tf.nn.relu, W_init=W_init, name='conv1_1')(ni)
+    nn = MaxPool2d((2, 2), (2, 2), padding='SAME', name='pool1_1')(nn)
+    nn = Conv2d(64, (3, 3), (1, 1), padding='SAME', act=tf.nn.relu, W_init=W_init, b_init=None, name='conv1_2')(nn)
+    nn = MaxPool2d((2, 2), (2, 2), padding='SAME', name='pool1_2')(nn)
+
+    nn = Conv2d(128, (3, 3), (1, 1), padding='SAME', act=tf.nn.relu, W_init=W_init, b_init=None, name='conv2_1')(nn)
+    nn = MaxPool2d((2, 2), (2, 2), padding='SAME', name='pool2_1')(nn)
+    nn = Conv2d(128, (3, 3), (1, 1), padding='SAME', act=tf.nn.relu, W_init=W_init, b_init=None, name='conv2_2')(nn)
+    nn = MaxPool2d((2, 2), (2, 2), padding='SAME', name='pool2_2')(nn)
+
+    nn = Conv2d(256, (3, 3), (1, 1), padding='SAME', act=tf.nn.relu, W_init=W_init, b_init=None, name='conv3_1')(nn)
+    nn = MaxPool2d((2, 2), (2, 2), padding='SAME', name='pool3_1')(nn)
+    nn = Conv2d(256, (3, 3), (1, 1), padding='SAME', act=tf.nn.relu, W_init=W_init, b_init=None, name='conv3_2')(nn)
+    nn = MaxPool2d((2, 2), (2, 2), padding='SAME', name='pool3_2')(nn)
+
+    nn = Conv2d(512, (3, 3), (1, 1), padding='SAME', act=tf.nn.relu, W_init=W_init, b_init=None, name='conv4_1')(nn)
+    nn = MaxPool2d((2, 2), (2, 2), padding='SAME', name='pool4_1')(nn)
+    nn = Conv2d(512, (3, 3), (1, 1), padding='SAME', act=tf.nn.relu, W_init=W_init, b_init=None, name='conv4_2')(nn)
+    nn = MaxPool2d((2, 2), (2, 2), padding='SAME', name='pool4_2')(nn)
+
+    nn = Flatten(name='flatten')(nn)
+    nn = Dense(1000, act=None, W_init=W_init2, name='output')(nn)
+
+    M = Model(inputs=ni, outputs=nn, name='cnn')
+    return M
+
+
+net = create_model([None, 224, 224, 3])
+# loaded weights whose name is not found in network's weights will be skipped.
+# If ckpt has the same naming rule as TL, We can restore the model with tl.files.load_and_assign_ckpt(model_dir=, network=, skip=True)
+tl.files.load_and_assign_npz_dict(network=net, skip=True)
+
+# you can use the following code to view the restore the model parameters.
+net_weights_name = [w.name for w in net.all_weights]
+for i in range(len(net_weights_name)):
+    print(net_weights_name[i], net.all_weights[net_weights_name.index(net_weights_name[i])])

tensorlayer/files/__init__.py

@@ -72,4 +72,6 @@
     #'load_graph',
     #'save_graph_and_params',
     #'load_graph_and_params',
+    'load_and_assign_ckpt',
+    'ckpt_to_npz_dict'
 ]

tensorlayer/files/utils.py

@@ -26,6 +26,7 @@
 from tensorflow.python.platform import gfile
 from tensorflow.python.util import serialization
 from tensorflow.python.util.tf_export import keras_export
+from tensorflow.python import pywrap_tensorflow
 
 import progressbar
 import tensorlayer as tl
@@ -76,6 +77,8 @@
     'static_graph2net',
     # 'save_pkl_graph',
     # 'load_pkl_graph',
+    'load_and_assign_ckpt',
+    'ckpt_to_npz_dict',
 ]
 
 
@@ -2775,3 +2778,119 @@ def load_hdf5_to_weights(filepath, network, skip=False):
 
     f.close()
     logging.info("[*] Load %s SUCCESS!" % filepath)
+
+
+def check_ckpt_file(model_dir):
+    model_dir = model_dir
+    model_path = None
+    count_extension = 0
+    for root, dirs, files in os.walk(model_dir):
+        for file in files:
+            filename, extension = os.path.splitext(file)
+            if extension in ['.data-00000-of-00001', '.index', '.meta']:
+                count_extension += 1
+        if count_extension == 3:
+            model_path = model_dir + '/' + filename
+        else:
+            raise Exception("Check the file extension for missing .data-00000-of-00001, .index, .meta")
+        if model_path is None:
+            raise Exception('The ckpt file is not found')
+    return model_path, filename
+
+
+def rename_weight_or_biases(variable_name):
+    if variable_name is None:
+        return variable_name
+    split_var = variable_name.split('/')
+
+    str_temp = ''
+    for i in range(len(split_var)):
+        if 'w' in split_var[i]:
+            split_var[i] = 'filters:0'
+        elif 'b' in split_var[i]:
+            split_var[i] = 'biases:0'
+        else:
+            pass
+
+        if i < len(split_var) - 1:
+            str_temp = str_temp + split_var[i] + '/'
+        else:
+            str_temp = str_temp + split_var[i]
+
+    return str_temp
+
+
+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):
+    """ Save ckpt weights to npz file
+
+    Parameters
+    ----------
+    model_dir : str
+        Filename to which the weights will be loaded, should be of ckpt format.
+        Examples: model_dir = /root/cnn_model/
+    save_name : str
+        The save_name of the `.npz` file.
+    rename_key : bool
+        Modify parameter naming,  used to match TL naming rule.
+        Examples: conv1_1/b_b --> conv1_1/biases:0 ; conv1_1/w_w --> conv1_1/filters:0
+
+    Returns
+    -------
+
+    """
+    model_path, _ = check_ckpt_file(model_dir)
+
+    reader = pywrap_tensorflow.NewCheckpointReader(model_path)
+    var_to_shape_map = reader.get_variable_to_shape_map()
+
+    parameters_dict = {}
+    if rename_key is False:
+        for key in sorted(var_to_shape_map):
+            parameters_dict[key] = reader.get_tensor(key)
+    elif rename_key is True:
+        for key in sorted(var_to_shape_map):
+            parameters_dict[rename_weight_or_biases(key)] = reader.get_tensor(key)
+
+    np.savez(save_name, **parameters_dict)
+    parameters_dict = None
+    del parameters_dict
+    logging.info("[*] Ckpt weights saved in npz_dict %s" % save_name)