update optimizer and dataflow

Author: hanjr92

examples/basic_tutorials/tutorial_tensorlayer_model_load.py

@@ -2,8 +2,9 @@
 # -*- 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
@@ -105,7 +106,7 @@ def forward(self, x):
 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)
+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.

tensorlayerx/backend/ops/mindspore_backend.py

@@ -1170,6 +1170,8 @@ def floor(x):
 
 def gather(params, indices, axis=None):
     op = P.Gather()
+    if axis is None:
+        axis = 0
     return op(params, indices, axis)
 
 
@@ -1590,10 +1592,7 @@ def reduce_std(x, axis=None, keepdims=False):
 
 
 def reduce_sum(x, axis=None, keepdims=False):
-    op = P.ReduceSum(keep_dims=keepdims)
-    if axis is None:
-        return op(x)
-    return op(x, axis=axis)
+    return msnp.sum(x, axis=axis, keepdims=keepdims)
 
 
 def reduce_variance(x, axis=None, keepdims=False):
@@ -1729,11 +1728,15 @@ def tanh(x):
 
 def any(x, axis=None, keepdims=False):
     op = P.ReduceAny(keep_dims=keepdims)
+    if axis is None:
+        return op(x)
     return op(x, axis)
 
 
 def all(x, axis=None, keepdims=False):
     op = P.ReduceAll(keep_dims=keepdims)
+    if axis is None:
+        return op(x)
     return op(x, axis)
 
 
@@ -1779,8 +1782,7 @@ def zeros_like(x, dtype=None):
 
 
 def squeeze(x, axis=None):
-    op = P.Squeeze(axis)
-    return op(x)
+    return msnp.squeeze(x, axis)
 
 
 def unsorted_segment_sum(x, segment_ids, num_segments):
@@ -1792,7 +1794,7 @@ def unsorted_segment_sum(x, segment_ids, num_segments):
 def unsorted_segment_mean(x, segment_ids, num_segments):
     segment_ids = ms.Tensor(segment_ids)
     op = P.UnsortedSegmentSum()
-    x_one =  msnp.ones_like(x, dtype=x.dtype)
+    x_one = msnp.ones_like(x, dtype=x.dtype)
     sum = op(x, segment_ids, num_segments)
     one = op(x_one, segment_ids, num_segments)
     return sum/one

tensorlayerx/backend/ops/tensorflow_backend.py

@@ -1249,7 +1249,7 @@ def gather(params, indices, axis=None):
     indices : indices
         The index Tensor. Must be one of the following types: int32, int64. The values must be in range [0, params.shape[axis]).
     axis : tensor.
-        Must be one of the following types: int32, int64. The axis in params to gather indices from.
+        Must be one of the following types: int32, int64. The axis in params to gather indices from. The default value is None, if None, the axis is 0.
 
     Returns
     -------

tensorlayerx/backend/ops/torch_backend.py

@@ -1012,7 +1012,9 @@ def floor(x):
     return torch.floor(x)
 
 
-def gather(params, indices, axis = 0):
+def gather(params, indices, axis = None):
+    if axis is None:
+        axis = 0
     if axis < 0:
         axis = len(params.shape) + axis
     if axis == 0:
@@ -1522,11 +1524,16 @@ def tanh(x):
 
 
 def any(x, axis=None, keepdims=False):
-    return torch.any(x, dim=axis, keepdim=keepdims)
-
+    if axis is not None:
+        return torch.any(x, dim=axis, keepdim=keepdims)
+    else:
+        return torch.any(x)
 
 def all(x, axis=None, keepdims=False):
-    return  torch.all(x, dim=axis, keepdim=keepdims)
+    if axis is not None:
+        return torch.all(x, dim=axis, keepdim=keepdims)
+    else:
+        return torch.all(x)
 
 
 def logical_and(x, y):

tensorlayerx/dataflow/utils.py

@@ -22,7 +22,9 @@ def default_convert(data):
             data = tf.convert_to_tensor(data)
         elif BACKEND == 'torch':
             import torch
+            device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
             data = torch.as_tensor(data)
+            data = data.to(device)
         elif BACKEND == 'paddle':
             import paddle
             data = paddle.to_tensor(data)
@@ -76,17 +78,22 @@ def default_collate_torch(batch):
     data = batch[0]
     data_type = type(data)
     import torch
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
     if isinstance(data, torch.Tensor):
         batch = torch.stack(batch, 0)
+        batch = batch.to(device)
         return batch
     elif isinstance(data, np.ndarray):
         batch = np.stack(batch, axis=0)
         batch = torch.as_tensor(batch)
+        batch = batch.to(device)
         return batch
     elif isinstance(data, numbers.Number):
         batch = torch.as_tensor(batch)
+        batch = batch.to(device)
         return batch
     elif isinstance(data, (str, bytes)):
+        batch = batch.to(device)
         return batch
     elif isinstance(data, collections.abc.Mapping):
         return {key: default_collate_torch([d[key] for d in batch]) for key in data}

tensorlayerx/model/core.py

@@ -429,19 +429,16 @@ def th_train(
         self, n_epoch, train_dataset, network, loss_fn, train_weights, optimizer, metrics, print_train_batch,
         print_freq, test_dataset
     ):
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        network.to(device)
         for epoch in range(n_epoch):
             start_time = time.time()
 
             train_loss, train_acc, n_iter = 0, 0, 0
             for X_batch, y_batch in train_dataset:
-                device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
                 network.set_train()
-                X_batch = X_batch.to(device)
-                y_batch = y_batch.to(device)
-                network.to(device)
                 output = network(X_batch)
                 loss = loss_fn(output, y_batch)
-
                 grads = optimizer.gradient(loss, train_weights)
                 optimizer.apply_gradients(zip(grads, train_weights))
 

tensorlayerx/model/utils.py

@@ -177,15 +177,13 @@ def __call__(self, data, label):
 class TrainOneStepWithTH(object):
 
     def __init__(self, net_with_loss, optimizer, train_weights):
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
         self.net_with_loss = net_with_loss
+        self.net_with_loss.to(device)
         self.optimizer = optimizer
         self.train_weights = train_weights
 
     def __call__(self, data, label):
-        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-        data = data.to(device)
-        label = label.to(device)
-        self.net_with_loss.to(device)
         loss = self.net_with_loss(data, label)
         grads = self.optimizer.gradient(loss, self.train_weights)
         self.optimizer.apply_gradients(zip(grads, self.train_weights))