update torch device allocation

Author: hanjr92

tensorlayerx/backend/ops/mindspore_backend.py

@@ -72,7 +72,7 @@ def get_tensor_shape(x):
 
 
 # initializers
-def zeros(shape, dtype=mstype.float32):
+def zeros(shape, dtype=mstype.float32, device = None):
     """
     Creates a tensor with all elements set to zero.
 
@@ -95,7 +95,7 @@ def zeros(shape, dtype=mstype.float32):
     return Tensor(arr, dtype=dtype)
 
 
-def ones(shape, dtype=mstype.float32):
+def ones(shape, dtype=mstype.float32, device = None):
     """
     Creates a tensor with all elements set to ones.
 
@@ -118,7 +118,7 @@ def ones(shape, dtype=mstype.float32):
     return Tensor(arr, dtype=dtype)
 
 
-def constant(value, dtype=mstype.float32, shape=None):
+def constant(value, dtype=mstype.float32, shape=None, device = None):
     """
     Creates a constant tensor from a tensor-like object.
 
@@ -413,7 +413,7 @@ def xavier_normal(shape, dtype, seed=None):
     return Tensor(arr, dtype=dtype)
 
 
-def Variable(initial_value, name, trainable=True):
+def Variable(initial_value, name, trainable=True, device = None):
     """
     Creates a new variable with value initial_value.
 
@@ -622,7 +622,7 @@ def concat(values, axis):
     return outputs
 
 
-def convert_to_tensor(value, dtype=None):
+def convert_to_tensor(value, dtype=None, device = None):
     """
     Converts the given value to a Tensor.
 

tensorlayerx/backend/ops/paddle_backend.py

@@ -45,7 +45,7 @@ def get_tensor_shape(x):
 
 
 # initializers
-def zeros(shape, dtype="float32"):
+def zeros(shape, dtype="float32", device = None):
     """
     Creates a tensor with all elements set to zero.
 
@@ -64,7 +64,7 @@ def zeros(shape, dtype="float32"):
     return pd.zeros(shape=shape, dtype=dtype)
 
 
-def ones(shape, dtype="float32"):
+def ones(shape, dtype="float32", device = None):
     """
     Creates a tensor with all elements set to ones.
 
@@ -83,7 +83,7 @@ def ones(shape, dtype="float32"):
     return pd.ones(shape=shape, dtype=dtype)
 
 
-def constant(value, dtype="float32", shape=None):
+def constant(value, dtype="float32", shape=None, device = None):
     """
     Creates a constant tensor from a tensor-like object.
 
@@ -241,7 +241,7 @@ def xavier_uniform(shape, dtype, seed=None):
     raise NotImplementedError
 
 
-def Variable(initial_value, name, trainable=None):
+def Variable(initial_value, name, trainable=None, device = None):
     """
     Creates a new variable with value initial_value.
 
@@ -425,7 +425,7 @@ def concat(values, axis=0):
     return pd.concat(values, axis)
 
 
-def convert_to_tensor(value, dtype=None):
+def convert_to_tensor(value, dtype=None, device = None):
     """
     Converts the given value to a Tensor.
 

tensorlayerx/backend/ops/tensorflow_backend.py

@@ -83,7 +83,7 @@ def get_tensor_shape(x):
 
 
 # initializers
-def zeros(shape, dtype='float32'):
+def zeros(shape, dtype='float32', device = None):
     """
     Creates a tensor with all elements set to zero.
 
@@ -93,6 +93,8 @@ def zeros(shape, dtype='float32'):
         a tuple of integers, or a 1-D Tensor of type int32.
     dtype : tensor or str
         The DType of an element in the resulting Tensor
+    device : str or None
+        create a tensor on 'cpu' or 'gpu', defautl is None.
 
     Returns
     -------
@@ -109,7 +111,7 @@ def zeros(shape, dtype='float32'):
     return tf.zeros(shape=shape, dtype=dtype_str(dtype))
 
 
-def ones(shape, dtype='float32'):
+def ones(shape, dtype='float32', device = None):
     """
     Creates a tensor with all elements set to ones.
 
@@ -119,6 +121,8 @@ def ones(shape, dtype='float32'):
         a tuple of integers, or a 1-D Tensor of type int32.
     dtype : tensor or str
         The DType of an element in the resulting Tensor
+    device : str or None
+        create a tensor on 'cpu' or 'gpu', defautl is None.
 
     Returns
     -------
@@ -135,7 +139,7 @@ def ones(shape, dtype='float32'):
     return tf.ones(shape=shape, dtype=dtype_str(dtype))
 
 
-def constant(value, dtype='float32', shape=None):
+def constant(value, dtype='float32', shape=None, device = None):
     """
     Creates a constant tensor from a tensor-like object.
 
@@ -147,6 +151,8 @@ def constant(value, dtype='float32', shape=None):
          The type of the elements of the resulting tensor.
     shape : tuple
         Optional dimensions of resulting tensor.
+    device : str or None
+        create a tensor on 'cpu' or 'gpu', defautl is None.
 
     Returns
     -------
@@ -345,7 +351,7 @@ def xavier_uniform(shape, dtype='float32', seed=None):
     return tf.initializers.glorot_uniform(seed)(shape=shape, dtype=dtype_str(dtype))
 
 
-def Variable(initial_value, name, trainable=True):
+def Variable(initial_value, name, trainable=True, device = None):
     """
     Creates a new variable with value initial_value.
 
@@ -355,6 +361,8 @@ def Variable(initial_value, name, trainable=True):
         A Tensor, or Python object convertible to a Tensor
     name : str
         Optional name for the variable. Defaults to 'Variable' and gets uniquified automatically.
+    device : str or None
+        create a tensor on 'cpu' or 'gpu', defautl is None.
     Returns
     -------
         Variable
@@ -591,7 +599,7 @@ def concat(values, axis):
     return tf.concat(values, axis)
 
 
-def convert_to_tensor(value, dtype=None):
+def convert_to_tensor(value, dtype=None, device = None):
     """
     Converts the given value to a Tensor.
 
@@ -601,6 +609,8 @@ def convert_to_tensor(value, dtype=None):
         An object whose type has a registered Tensor conversion function.
     dtype : optional
         Optional element type for the returned tensor. If missing, the type is inferred from the type of value.
+    device : str or None
+        create a tensor on 'cpu' or 'gpu', defautl is None.
 
     Returns
     -------

tensorlayerx/backend/ops/torch_backend.py

@@ -52,7 +52,7 @@ def get_tensor_shape(x):
 
 
 # initializers
-def zeros(shape, dtype=None):
+def zeros(shape, dtype=None, device = None):
     """
     Creates a tensor with all elements set to zero.
 
@@ -68,11 +68,14 @@ def zeros(shape, dtype=None):
         A Tensor with all elements set to zero.
 
     """
+    if device == 'cpu':
+        device = torch.device('cpu')
+    elif device == 'gpu':
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    return torch.zeros(size=shape, dtype=dtype, device = device)
 
-    return torch.zeros(size=shape, dtype=dtype)
 
-
-def ones(shape, dtype=None):
+def ones(shape, dtype=None, device = None):
     """
     Creates a tensor with all elements set to ones.
 
@@ -88,11 +91,14 @@ def ones(shape, dtype=None):
         A Tensor with all elements set to zero.
 
     """
-
-    return torch.ones(size=shape, dtype=dtype)
+    if device == 'cpu':
+        device = torch.device('cpu')
+    elif device == 'gpu':
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    return torch.ones(size=shape, dtype=dtype, device = device)
 
 
-def constant(value, dtype=None, shape=None):
+def constant(value, dtype=None, shape=None, device =None):
     """
     Creates a constant tensor from a tensor-like object.
 
@@ -110,8 +116,11 @@ def constant(value, dtype=None, shape=None):
         A Constant Tensor.
 
     """
-
-    w = torch.empty(size=shape, dtype=dtype)
+    if device == 'cpu':
+        device = torch.device('cpu')
+    elif device == 'gpu':
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    w = torch.empty(size=shape, dtype=dtype, device = device)
     return torch.nn.init.constant_(w, value)
 
 
@@ -426,7 +435,7 @@ def concat(values, axis=0):
     return torch.cat(values, axis)
 
 
-def convert_to_tensor(value, dtype=None):
+def convert_to_tensor(value, dtype=None, device = None):
     """
     Converts the given value to a Tensor.
 
@@ -443,14 +452,18 @@ def convert_to_tensor(value, dtype=None):
     """
     if isinstance(dtype, str):
         dtype = _dtypeDict[dtype]
-    return torch.tensor(value, dtype=dtype)
+    if device == 'cpu':
+        device = torch.device('cpu')
+    elif device == 'gpu':
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    return torch.tensor(value, dtype=dtype, device = device)
 
 
 def convert_to_numpy(value):
     try:
         return value.numpy()
     except:
-        return value.detach().numpy()
+        return value.cpu().detach().numpy()
 
 
 def sqrt(x):
@@ -1505,7 +1518,7 @@ def tanh(x):
         A Tensor. Has the same type as x.
     """
 
-    return F.tanh(x)
+    return torch.tanh(x)
 
 
 def any(x, axis=None, keepdims=False):

tensorlayerx/dataflow/dataloader.py

@@ -61,14 +61,11 @@ def __init__(
         batch_sampler=None,
         num_workers=0,
         collate_fn=None,
-        # pin_memory = False,
         time_out=0,
         worker_init_fn=None,
-        #multiprocessing_context=None,
         prefetch_factor=2,
         persistent_workers=False,
     ):
-        # assert isinstance(dataset, Dataset), "dataset should be subclass of tensorlayerx.dataflow.Dataset"
         self.dataset = dataset
         assert num_workers >= 0, "num_workers should be a non_negative integer"
         if num_workers == 0 and prefetch_factor != 2:
@@ -77,10 +74,8 @@ def __init__(
             raise ValueError('persistent_workers option needs num_workers > 0')
         self.num_workers = num_workers
         self.prefetch_factor = prefetch_factor
-        # self.pin_memory = pin_memory
         self.time_out = time_out
         self.worker_init_fn = worker_init_fn
-        #self.multiprocessing_context = multiprocessing_context
         if isinstance(dataset, IterableDataset):
             self._dataset_kind = _DatasetKind.Iter
             if shuffle is not False:

tensorlayerx/metrics/torch_metric.py

@@ -51,7 +51,7 @@ def update(self, y_pred, y_true):
             y_true = torch.argmax(y_true, dim=-1, keepdim=True)
         correct = y_pred == y_true
         correct = correct.to(torch.float32)
-        correct = correct.numpy()
+        correct = correct.cpu().numpy()
         num_samples = np.prod(np.array(correct.shape[:-1]))
         num_corrects = correct[..., :self.topk].sum()
         self.total = num_corrects
@@ -78,12 +78,12 @@ def __init__(
 
     def update(self, y_pred, y_true):
         if isinstance(y_true, torch.Tensor):
-            y_true = y_true.numpy()
+            y_true = y_true.cpu().numpy()
         elif not isinstance(y_pred, np.ndarray):
             raise TypeError("The y_true must be a numpy array or Tensor.")
 
         if isinstance(y_pred, torch.Tensor):
-            y_pred = y_pred.numpy()
+            y_pred = y_pred.cpu().numpy()
         elif not isinstance(y_pred, np.ndarray):
             raise TypeError("The y_pred must be a numpy array or Tensor.")
 
@@ -131,12 +131,12 @@ def __init__(self):
 
     def update(self, y_pred, y_true):
         if isinstance(y_true, torch.Tensor):
-            y_true = y_true.numpy()
+            y_true = y_true.cpu().numpy()
         elif not isinstance(y_pred, np.ndarray):
             raise TypeError("The y_true must be a numpy array or Tensor.")
 
         if isinstance(y_pred, torch.Tensor):
-            y_pred = y_pred.numpy()
+            y_pred = y_pred.cpu().numpy()
         elif not isinstance(y_pred, np.ndarray):
             raise TypeError("The y_pred must be a numpy array or Tensor.")
 
@@ -169,12 +169,12 @@ def __init__(self):
 
     def update(self, y_pred, y_true):
         if isinstance(y_true, torch.Tensor):
-            y_true = y_true.numpy()
+            y_true = y_true.cpu().numpy()
         elif not isinstance(y_pred, np.ndarray):
             raise TypeError("The y_true must be a numpy array or Tensor.")
 
         if isinstance(y_pred, torch.Tensor):
-            y_pred = y_pred.numpy()
+            y_pred = y_pred.cpu().numpy()
         elif not isinstance(y_pred, np.ndarray):
             raise TypeError("The y_pred must be a numpy array or Tensor.")
 
@@ -209,7 +209,7 @@ def acc(predicts, labels, topk=1):
         y_true = torch.argmax(labels, dim=-1, keepdim=True)
     correct = y_pred == y_true
     correct = correct.to(torch.float32)
-    correct = correct.numpy()
+    correct = correct.cpu().numpy()
     num_samples = np.prod(np.array(correct.shape[:-1]))
     num_corrects = correct[..., :topk].sum()
     total = num_corrects

tensorlayerx/model/core.py

@@ -59,10 +59,10 @@ class Model:
     >>> class Net(Module):
     >>>     def __init__(self):
     >>>         super(Net, self).__init__()
-    >>>         self.conv = tlx.nn.Conv2d(n_filter=32, filter_size=(3, 3), strides=(2, 2), in_channels=5, name='conv2d')
+    >>>         self.conv = tlx.nn.Conv2d(out_channels=32, kernel_size=(3, 3), stride=(2, 2), in_channels=5, name='conv2d')
     >>>         self.bn = tlx.nn.BatchNorm2d(num_features=32, act=tlx.ReLU)
     >>>         self.flatten = tlx.nn.Flatten()
-    >>>         self.fc = tlx.nn.Dense(n_units=12, in_channels=32*224*224) # padding=0
+    >>>         self.fc = tlx.nn.Linear(out_features=12, in_features=32*224*224) # padding=0
     >>>
     >>>     def construct(self, x):
     >>>         x = self.conv(x)
@@ -434,8 +434,11 @@ def th_train(
 
             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)