update ops and parameter container

Author: hanjr92

docs/modules/nn.rst

@@ -16,6 +16,9 @@ Layer list
    Sequential
    ModuleList
    ModuleDict
+   Parameter
+   ParameterList
+   ParameterDict
 
    Input
 
@@ -139,6 +142,18 @@ ModuleDict
 ^^^^^^^^^^^^^^^^
 .. autoclass:: ModuleDict
 
+Parameter
+^^^^^^^^^^^^^^^^
+.. autofunction:: Parameter
+
+ParameterList
+^^^^^^^^^^^^^^^^
+.. autoclass:: ParameterList
+
+ParameterDict
+^^^^^^^^^^^^^^^^
+.. autoclass:: ParameterDict
+
 .. -----------------------------------------------------------
 ..                        Input Layer
 .. -----------------------------------------------------------

docs/modules/ops.rst

@@ -117,6 +117,7 @@ API - Operations
    mask_select
    eye
    einsum
+   set_device
 
 TensorLayerX Tensor Operations
 --------------------------------
@@ -563,4 +564,8 @@ eye
 
 einsum
 ^^^^^^^^^^^^^^^^^^^^^^^
-.. autofunction:: einsum
+.. autofunction:: einsum
+
+set_device
+^^^^^^^^^^^^^^^^^^^^^^^
+.. autofunction:: set_device

examples/basic_tutorials/tutorial_ModuleContainer.py renamed to examples/basic_tutorials/Module_Container.py

@@ -0,0 +1,33 @@
+import os
+os.environ['TL_BACKEND'] = 'tensorflow'
+# os.environ['TL_BACKEND'] = 'mindspore'
+# os.environ['TL_BACKEND'] = 'paddle'
+# os.environ['TL_BACKEND'] = 'torch'
+
+import tensorlayerx as tlx
+from tensorlayerx.nn import Module, Parameter, ParameterList, ParameterDict
+tlx.set_device(device='CPU', id = 0)
+
+class MyModule(Module):
+    def __init__(self):
+        super(MyModule, self).__init__()
+        self.params1 = ParameterDict({
+                'left': Parameter(tlx.ones((5, 10))),
+                'right': Parameter(tlx.zeros((5, 10)))
+        })
+
+        self.params2 = ParameterList(
+            [Parameter(tlx.ones((10,5))), Parameter(tlx.ones((5,10)))]
+        )
+
+    def forward(self, x, choice):
+        x = tlx.matmul(x, self.params1[choice])
+        x = tlx.matmul(x, self.params2[0])
+        x = tlx.matmul(x, self.params2[1])
+        return x
+
+input = tlx.nn.Input(shape=(5,5))
+net = MyModule()
+
+output = net(input, choice = 'right')
+print(output)

examples/basic_tutorials/Parameter_Container.py

@@ -147,6 +147,8 @@
 from .ops import diag
 from .ops import mask_select
 from .ops import eye
+from .ops import einsum
+from .ops import set_device
 # dtype
 from .ops import (
     DType, float16, float32, float64, int8, int16, int32, int64, uint8, uint16, uint32, uint64, bool, complex64,

tensorlayerx/backend/__init__.py

@@ -84,7 +84,6 @@
 from .load_backend import QuanConvBn
 
 # load ops
-from .load_backend import einsum
 from .load_backend import Variable
 from .load_backend import matmul
 from .load_backend import add
@@ -200,6 +199,8 @@
 from .load_backend import diag
 from .load_backend import mask_select
 from .load_backend import eye
+from .load_backend import einsum
+from .load_backend import set_device
 # dtype
 from .load_backend import (
     DType, float16, float32, float64, int8, int16, int32, int64, uint8, uint16, uint32, uint64, bool, complex64,

tensorlayerx/backend/ops/__init__.py

@@ -58,7 +58,7 @@
     import mindspore.context as context
     import os
     os.environ['DEVICE_ID'] = '0'
-    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU'),
+    context.set_context(mode=context.PYNATIVE_MODE),
     # context.set_context(mode=context.PYNATIVE_MODE, device_target='CPU'),
     # enable_task_sink=True, enable_loop_sink=True)
     # context.set_context(mode=context.PYNATIVE_MODE, device_target='Ascend')

tensorlayerx/backend/ops/load_backend.py

@@ -1838,4 +1838,10 @@ def __init__(self, equation):
         self.einsum = ms.ops.Einsum(equation)
 
     def __call__(self, *args):
-        return self.einsum(tuple(args))
+        return self.einsum(tuple(args))
+
+def set_device(device = 'GPU', id = 0):
+    if device not in ['GPU', 'CPU', 'Ascend']:
+        raise ValueError ("In mindspore, only support 'CPU', 'GPU' and 'Ascend'.")
+    ms.context.set_context(device_target=device)
+    ms.context.set_context(device_id = id)

tensorlayerx/backend/ops/mindspore_backend.py

@@ -1856,3 +1856,9 @@ def __init__(self, equation):
     def __call__(self, *args):
         return einsum(self.equation, *args)
 
+
+def set_device(device = 'GPU', id = 0):
+    device = device.lower()
+    if device == 'GPU':
+        device = device + ':' + str(id)
+    paddle.device.set_device(device)

tensorlayerx/backend/ops/paddle_backend.py

@@ -3867,4 +3867,27 @@ def __init__(self, equation):
         self.equation = equation
 
     def __call__(self, *args):
-        return tf.einsum(self.equation, *args)
+        return tf.einsum(self.equation, *args)
+
+def set_device(device = 'GPU', id = 0):
+    """This function can specify the global device which the OP will run.
+
+    Parameters
+    ----------
+    device : str
+        Specific running device. It can be 'CPU', 'GPU' and 'Ascend'(In mindspore backend).
+    id : int
+        Device id.
+
+    """
+    if device not in ['GPU', 'CPU']:
+        raise ValueError ("Only support 'CPU', 'GPU'.")
+    if device == 'GPU':
+        gpus = tf.config.experimental.list_physical_devices('GPU')
+        if gpus:
+            try:
+                for gpu in gpus:
+                    tf.config.experimental.set_memory_growth(gpu, True)
+                tf.config.experimental.set_visible_devices(gpus[id], 'GPU')
+            except RuntimeError as e:
+                print(e)