add tutorial load ckpt

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/utils.py

@@ -2780,6 +2780,46 @@ def load_hdf5_to_weights(filepath, network, skip=False):
     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
 
@@ -2798,16 +2838,7 @@ def load_and_assign_ckpt(model_dir, network=None, skip=True):
     -------
 
     """
-    model_dir = model_dir
-    model_path = None
-    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']:
-                model_path = model_dir + '/' + filename
-                break
-        if model_path == None:
-            raise Exception('The ckpt file is not found')
+    model_path, filename = check_ckpt_file(model_dir)
 
     reader = pywrap_tensorflow.NewCheckpointReader(model_path)
     var_to_shape_map = reader.get_variable_to_shape_map()
@@ -2828,7 +2859,7 @@ def load_and_assign_ckpt(model_dir, network=None, skip=True):
     logging.info("[*] Model restored from ckpt %s" % filename)
 
 
-def ckpt_to_npz_dict(model_dir, save_name='model.npz'):
+def ckpt_to_npz_dict(model_dir, save_name='model.npz', rename_key=False):
     """ Save ckpt weights to npz file
 
     Parameters
@@ -2838,28 +2869,27 @@ def ckpt_to_npz_dict(model_dir, save_name='model.npz'):
         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_dir = model_dir
-    model_path = None
-    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']:
-                model_path = model_dir + '/' + filename
-                break
-        if model_path == None:
-            raise Exception('The ckpt file is not found')
+    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 = {}
-    for key in sorted(var_to_shape_map):
-        parameters_dict[key] = reader.get_tensor(key)
+    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