Update TLX training models can be imported into any backend. (#17)

Author: Laicheng0830

examples/basic_tutorials/tutorial_tensorlayer_model_load.py

@@ -2,33 +2,14 @@
 # -*- coding: utf-8 -*-
 
 import os
-# os.environ['TL_BACKEND'] = 'tensorflow'
-os.environ['TL_BACKEND'] = 'paddle'
+os.environ['TL_BACKEND'] = 'tensorflow'
+# os.environ['TL_BACKEND'] = 'paddle'
 # os.environ['TL_BACKEND'] = 'mindspore'
 # os.environ['TL_BACKEND'] = 'torch'
 
 import tensorlayerx as tlx
 from tensorlayerx.nn import Module
 from tensorlayerx.nn import Linear, Dropout, Conv2d, MaxPool2d, Flatten
-from tensorlayerx.dataflow import Dataset
-
-X_train, y_train, X_val, y_val, X_test, y_test = tlx.files.load_mnist_dataset(shape=(-1, 784))
-
-
-class mnistdataset(Dataset):
-
-    def __init__(self, data=X_train, label=y_train):
-        self.data = data
-        self.label = label
-
-    def __getitem__(self, index):
-        data = self.data[index].astype('float32')
-        label = self.label[index].astype('int64')
-        return data, label
-
-    def __len__(self):
-        return len(self.data)
-
 
 class CustomModel(Module):
 
@@ -92,26 +73,23 @@ def forward(self, x):
         return z
 
 
-# TODO The MLP model was saved to the standard npz_dict format after training at the TensorFlow backend
-#  and imported into TensorFlow/PyTorch/PaddlePaddle/MindSpore.
+# # TODO The MLP model was saved to the standard npz_dict format after training at the TensorFlow backend
+# #  and imported into TensorFlow/PyTorch/PaddlePaddle/MindSpore.
 # MLP = CustomModel()
-# MLP.save_standard_weights('./model.npz')
-# # MLP.load_standard_weights('./model.npz', skip=True)
+# # MLP.save_standard_weights('./model.npz')
+# MLP.load_standard_weights('./model.npz', weights_from='tensorflow', weights_to='mindspore')
 # MLP.set_eval()
 # inputs = tlx.layers.Input(shape=(10, 784))
-# print(MLP(inputs))
+# output = MLP(inputs)
+# print(output)
 
 # TODO The CNN model was saved to the standard npz_dict format after training at the TensorFlow backend
 #  and imported into TensorFlow/PyTorch/PaddlePaddle/MindSpore.
 cnn = CNN()
-# cnn.save_standard_weights('./model.npz')
-# TODO Tensorflow trained parameters are imported to the TensorFlow backend.
-cnn.load_standard_weights('./model.npz', skip=False, reshape=True)
-
-# TODO Tensorflow backend trained parameters imported to PaddlePaddle/PyTorch/MindSpore to
-#  set reshape to True parameter to convert convolution shape.
-# cnn.load_standard_weights('./model.npz', skip=True, reshape=True)
+# cnn.save_standard_weights('./cnn.npz')
+cnn.load_standard_weights('./cnn.npz', weights_from='torch', weights_to='tensorflow')
 cnn.set_eval()
+
 inputs = tlx.nn.Input(shape=(10, 28, 28, 3), dtype=tlx.float32)
 outputs = cnn(inputs)
-print(outputs)
+# print(outputs)

tensorlayerx/nn/core/common.py

@@ -268,38 +268,20 @@ def _save_standard_weights_dict(net, file_path):
 
 def encode_list_name(list_name):
     # TensorFlow weights format: conv1.weight:0, conv1.bias:0
-    # Paddle weights format: conv1.weight, conv1.bias
-    # PyTorch weights format: conv1.W, conv1.W
+    # Paddle weights format: conv1.weights, conv1.bias
+    # PyTorch weights format: conv1.weights, conv1.bias
     # MindSpore weights format: conv1.weights, conv1.bias
     # standard weights format: conv1.weights, conv1.bias
 
     for i in range(len(list_name)):
         if tlx.BACKEND == 'tensorflow':
             list_name[i] = list_name[i][:-2]
-        if tlx.BACKEND == 'torch':
-            if list_name[i][-1] == 'W' and 'conv' not in list_name[i]:
-                list_name[i] = list_name[i][:-2] + str('/weights')
-            elif list_name[i][-1] == 'W' and 'conv' in list_name[i]:
-                list_name[i] = list_name[i][:-2] + str('/filters')
-            elif list_name[i][-1] == 'b':
-                list_name[i] = list_name[i][:-2] + str('/biases')
-            elif list_name[i].split('.')[-1] in ['beta', 'gamma', 'moving_mean', 'moving_var']:
-                pass
-            else:
-                raise NotImplementedError('This weights cannot be converted.')
     return list_name
 
 
 def decode_key_name(key_name):
     if tlx.BACKEND == 'tensorflow':
         key_name = key_name + str(':0')
-    if tlx.BACKEND == 'torch':
-        if key_name.split('/')[-1] in ['weights', 'filters']:
-            key_name = key_name[:-8] + str('.W')
-        elif key_name.split('/')[-1] == 'biases':
-            key_name = key_name[:-7] + str('.b')
-        else:
-            raise NotImplementedError('This weights cannot be converted.')
     return key_name
 
 
@@ -347,11 +329,30 @@ def save_standard_npz_dict(save_list=None, name='model.npz'):
     logging.info("[*] Model saved in npz_dict %s" % name)
 
 
-def _load_standard_weights_dict(net, file_path, skip=False, reshape=False, format='npz_dict'):
-    if format == 'npz_dict':
-        load_and_assign_standard_npz_dict(net, file_path, skip, reshape)
-    elif format == 'npz':
-        load_and_assign_standard_npz(file_path, net, reshape)
+def _load_standard_weights_dict(net, file_path, skip=False, weights_from='tensorflow', weights_to='tensorflow'):
+    """
+
+    Parameters
+    ----------
+    file_path : str
+        Name of the saved file.
+    skip : boolean
+        If 'skip' == True, loaded layer whose name is not found in 'layers' will be skipped. If 'skip' is False,
+        error will be raised when mismatch is found. Default False.
+    weights_from : string
+        The weights file is saved by which framework training. It has to be one of tensorflow,mindspore,paddle or torch.
+    weights_to : string
+        Which framework the weights file imports.It has to be one of tensorflow,mindspore,paddle or torch.
+    """
+    if weights_from == weights_to:
+        reshape = False
+    if weights_from == 'tensorflow' and weights_to != 'tensorflow':
+        reshape = True
+    if weights_from != 'tensorflow' and weights_to == 'tensorflow':
+        reshape = True
+    if weights_from !='tensorflow' and weights_to != 'tensorflow':
+        reshape = False
+    load_and_assign_standard_npz_dict(net, file_path, skip, reshape)
 
 
 def load_and_assign_standard_npz_dict(net, file_path, skip=False, reshape=False):
@@ -382,101 +383,96 @@ def load_and_assign_standard_npz_dict(net, file_path, skip=False, reshape=False)
         else:
             if tlx.BACKEND == 'tensorflow':
                 reshape_weights = weight_reshape(weights[key], reshape)
-                check_reshape(reshape_weights, net.all_weights[net_weights_name.index(de_key)])
+                # check_reshape(reshape_weights, net.all_weights[net_weights_name.index(de_key)])
                 utils.assign_tf_variable(net.all_weights[net_weights_name.index(de_key)], reshape_weights)
             elif tlx.BACKEND == 'mindspore':
                 reshape_weights = weight_reshape(weights[key], reshape)
-                import mindspore as ms
                 assign_param = ms.Tensor(reshape_weights, dtype=ms.float32)
-                check_reshape(assign_param, net.all_weights[net_weights_name.index(de_key)])
+                # check_reshape(assign_param, net.all_weights[net_weights_name.index(de_key)])
                 utils.assign_ms_variable(net.all_weights[net_weights_name.index(de_key)], assign_param)
             elif tlx.BACKEND == 'paddle':
                 reshape_weights = weight_reshape(weights[key], reshape)
-                check_reshape(reshape_weights, net.all_weights[net_weights_name.index(de_key)])
+                # check_reshape(reshape_weights, net.all_weights[net_weights_name.index(de_key)])
                 utils.assign_pd_variable(net.all_weights[net_weights_name.index(de_key)], reshape_weights)
             elif tlx.BACKEND == 'torch':
                 reshape_weights = weight_reshape(weights[key], reshape)
-                check_reshape(reshape_weights, net.all_weights[net_weights_name.index(de_key)])
+                # check_reshape(reshape_weights, net.all_weights[net_weights_name.index(de_key)])
                 utils.assign_th_variable(torch_weights_dict[de_key], reshape_weights)
             else:
                 raise NotImplementedError('Not implemented')
 
     logging.info("[*] Model restored from npz_dict %s" % file_path)
 
 
-def load_and_assign_standard_npz(file_path=None, network=None, reshape=False):
-    if network is None:
-        raise ValueError("network is None.")
-
-    if not os.path.exists(file_path):
-        logging.error("file {} doesn't exist.".format(file_path))
-        return False
-    else:
-        weights = utils.load_npz(name=file_path)
-        ops = []
-        if tlx.BACKEND == 'tensorflow':
-            for idx, param in enumerate(weights):
-                param = weight_reshape(param, reshape)
-                check_reshape(param, network.all_weights[idx])
-                ops.append(network.all_weights[idx].assign(param))
-
-        elif tlx.BACKEND == 'mindspore':
-
-            class Assign_net(Cell):
-
-                def __init__(self, y):
-                    super(Assign_net, self).__init__()
-                    self.y = y
-
-                def construct(self, x):
-                    Assign()(self.y, x)
-
-            for idx, param in enumerate(weights):
-                assign_param = Tensor(param, dtype=ms.float32)
-                assign_param = weight_reshape(assign_param, reshape)
-                check_reshape(assign_param, network.all_weights[idx])
-                Assign()(network.all_weights[idx], assign_param)
-
-        elif tlx.BACKEND == 'paddle':
-            for idx, param in enumerate(weights):
-                param = weight_reshape(param, reshape)
-                check_reshape(param, network.all_weights[idx])
-                utils.assign_pd_variable(network.all_weights[idx], param)
-
-        elif tlx.BACKEND == 'torch':
-            for idx, param in enumerate(weights):
-                param = weight_reshape(param, reshape)
-                check_reshape(param, network.all_weights[idx])
-                utils.assign_th_variable(network.all_weights[idx], param)
-        else:
-            raise NotImplementedError("This backend is not supported")
-        return ops
-
-    logging.info("[*] Load {} SUCCESS!".format(file_path))
-
-
-def check_reshape(weight, shape_weights):
-    if len(weight.shape) >= 4 and weight.shape[::-1] == tuple(shape_weights.shape):
-        if tlx.BACKEND == 'tensorflow':
-
-            raise Warning(
-                'Set reshape to True only when importing weights from MindSpore/PyTorch/PaddlePaddle to TensorFlow.'
-            )
-        if tlx.BACKEND == 'torch':
-            raise Warning('Set reshape to True only when importing weights from TensorFlow to PyTorch.')
-        if tlx.BACKEND == 'paddle':
-            raise Warning('Set reshape to True only when importing weights from TensorFlow to PaddlePaddle.')
-        if tlx.BACKEND == 'mindspore':
-            raise Warning('Set reshape to True only when importing weights from TensorFlow to MindSpore.')
+# def load_and_assign_standard_npz(file_path=None, network=None, reshape=False):
+#     if network is None:
+#         raise ValueError("network is None.")
+#
+#     if not os.path.exists(file_path):
+#         logging.error("file {} doesn't exist.".format(file_path))
+#         return False
+#     else:
+#         weights = utils.load_npz(name=file_path)
+#         ops = []
+#         if tlx.BACKEND == 'tensorflow':
+#             for idx, param in enumerate(weights):
+#                 param = weight_reshape(param, reshape)
+#                 check_reshape(param, network.all_weights[idx])
+#                 ops.append(network.all_weights[idx].assign(param))
+#
+#         elif tlx.BACKEND == 'mindspore':
+#             for idx, param in enumerate(weights):
+#                 assign_param = Tensor(param, dtype=ms.float32)
+#                 assign_param = weight_reshape(assign_param, reshape)
+#                 check_reshape(assign_param, network.all_weights[idx])
+#                 utils.assign_ms_variable(network.all_weights[idx], assign_param)
+#
+#         elif tlx.BACKEND == 'paddle':
+#             for idx, param in enumerate(weights):
+#                 param = weight_reshape(param, reshape)
+#                 check_reshape(param, network.all_weights[idx])
+#                 utils.assign_pd_variable(network.all_weights[idx], param)
+#
+#         elif tlx.BACKEND == 'torch':
+#             for idx, param in enumerate(weights):
+#                 param = weight_reshape(param, reshape)
+#                 check_reshape(param, network.all_weights[idx])
+#                 utils.assign_th_variable(network.all_weights[idx], param)
+#         else:
+#             raise NotImplementedError("This backend is not supported")
+#         return ops
+#
+#     logging.info("[*] Load {} SUCCESS!".format(file_path))
+
+
+# def check_reshape(weight, shape_weights):
+#     if len(weight.shape) >= 4 and weight.shape[::-1] == tuple(shape_weights.shape):
+#         if tlx.BACKEND == 'tensorflow':
+#
+#             raise Warning(
+#                 'Set reshape to True only when importing weights from MindSpore/PyTorch/PaddlePaddle to TensorFlow.'
+#             )
+#         if tlx.BACKEND == 'torch':
+#             raise Warning('Set reshape to True only when importing weights from TensorFlow to PyTorch.')
+#         if tlx.BACKEND == 'paddle':
+#             raise Warning('Set reshape to True only when importing weights from TensorFlow to PaddlePaddle.')
+#         if tlx.BACKEND == 'mindspore':
+#             raise Warning('Set reshape to True only when importing weights from TensorFlow to MindSpore.')
 
 
 def weight_reshape(weight, reshape=False):
     # TODO In this case only 2D convolution is considered. 3D convolution tests need to be supplemented.
     if reshape:
         if len(weight.shape) == 4:
-            weight = np.moveaxis(weight, (2, 3), (1, 0))
+            if tlx.BACKEND == 'tensorflow':
+                weight = np.moveaxis(weight, (1, 0), (2, 3))
+            else:
+                weight = np.moveaxis(weight, (2, 3), (1, 0))
         if len(weight.shape) == 5:
-            weight = np.moveaxis(weight, (3, 4), (1, 0))
+            if tlx.BACKEND == 'tensorflow':
+                weight = np.moveaxis(weight, (1, 0), (3, 4))
+            else:
+                weight = np.moveaxis(weight, (3, 4), (1, 0))
     return weight
 
 def tolist(tensors):

tensorlayerx/nn/core/core_mindspore.py

@@ -117,8 +117,8 @@ def load_weights(self, file_path, format=None, in_order=True, skip=False):
     def save_standard_weights(self, file_path):
         _save_standard_weights_dict(self, file_path)
 
-    def load_standard_weights(self, file_path, skip=False, reshape=False, format='npz_dict'):
-        _load_standard_weights_dict(self, file_path, skip, reshape, format)
+    def load_standard_weights(self, file_path, weights_from, weights_to, skip=False):
+        _load_standard_weights_dict(self, file_path, skip=skip, weights_from=weights_from, weights_to=weights_to)
 
     @staticmethod
     def _compute_shape(tensors):
@@ -158,6 +158,8 @@ def set_eval(self):
         """
         self._phase = 'predict'
         self.add_flags_recursive(training=False)
+        for layer in self.cells():
+            layer.is_train = False
         return self
 
     def test(self):
@@ -760,7 +762,6 @@ def update(self, parameters):
                         "ParameterDict update sequence element "
                         "#" + str(j) + " should be Iterable; is" + type(p).__name__
                     )
-                print(p)
                 if not len(p) == 2:
                     raise ValueError(
                         "ParameterDict update sequence element "

tensorlayerx/nn/core/core_paddle.py

@@ -270,8 +270,8 @@ def load_weights(self, file_path, format=None, in_order=True, skip=False):
     def save_standard_weights(self, file_path):
         _save_standard_weights_dict(self, file_path)
 
-    def load_standard_weights(self, file_path, skip=False, reshape=False, format='npz_dict'):
-        _load_standard_weights_dict(self, file_path, skip, reshape, format)
+    def load_standard_weights(self, file_path, weights_from, weights_to, skip=False):
+        _load_standard_weights_dict(self, file_path, skip=skip, weights_from=weights_from, weights_to=weights_to)
 
     def str_to_init(self, initializer):
         return str2init(initializer)
@@ -730,7 +730,6 @@ def update(self, parameters):
                         "ParameterDict update sequence element "
                         "#" + str(j) + " should be Iterable; is" + type(p).__name__
                     )
-                print(p)
                 if not len(p) == 2:
                     raise ValueError(
                         "ParameterDict update sequence element "

tensorlayerx/nn/core/core_tensorflow.py

@@ -221,7 +221,7 @@ def save_standard_weights(self, file_path):
 
         _save_standard_weights_dict(self, file_path)
 
-    def load_standard_weights(self, file_path, skip=False, reshape=False, format='npz_dict'):
+    def load_standard_weights(self, file_path, weights_from, weights_to, skip=False):
         """
 
         Parameters
@@ -231,14 +231,13 @@ def load_standard_weights(self, file_path, skip=False, reshape=False, format='np
         skip : boolean
             If 'skip' == True, loaded layer whose name is not found in 'layers' will be skipped. If 'skip' is False,
             error will be raised when mismatch is found. Default False.
-        reshape : boolean
-            This parameter needs to be set to True when importing parameters from tensorflow training to paddle/mindspore/pytorch,
-            and similarly when importing parameters from paddle/mindspore/pytorch training to tensorflow.
-            This parameter does not need to be set between paddle/mindspore/pytorch.
-
+        weights_from : string
+            The weights file is saved by which framework training. It has to be one of tensorflow,mindspore,paddle or torch.
+        weights_to : string
+            Which framework the weights file imports.It has to be one of tensorflow,mindspore,paddle or torch.
         """
 
-        _load_standard_weights_dict(self, file_path, skip, reshape, format)
+        _load_standard_weights_dict(self, file_path, skip=skip, weights_from=weights_from, weights_to=weights_to)
 
     def _set_mode_for_layers(self, is_train):
         """Set all layers of this network to a given mode.