Compatible with MindSpore1.6.1 and multiple back-end Computation Graph

Author: Laicheng0830

requirements/requirements_ms.txt

@@ -1 +1 @@
-mindspore==1.5.2
+mindspore==1.6.1

tensorlayerx/backend/ops/load_backend.py

@@ -59,7 +59,7 @@
     import os
     os.environ['DEVICE_ID'] = '0'
     context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU'),
-    # context.set_context(mode=context.GRAPH_MODE, device_target='CPU'),
+    # 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')
     sys.stderr.write('Using MindSpore backend.\n')

tensorlayerx/backend/ops/mindspore_nn.py

@@ -19,7 +19,7 @@
 from mindspore.ops.operations import LayerNorm
 import mindspore.numpy as np
 from mindspore.common.parameter import ParameterTuple
-from mindspore.nn.layer.rnns import _DynamicRNN
+from mindspore.nn.layer.rnns import _DynamicRNNBase
 import warnings
 import math
 
@@ -1953,7 +1953,7 @@ def __init__(
         self.w_hh_list = ParameterTuple(w_hh)
         self.b_ih_list = ParameterTuple(b_ih)
         self.b_hh_list = ParameterTuple(b_hh)
-        self.rnn = _DynamicRNN(mode)
+        self.rnn = _DynamicRNNBase(mode)
         self.is_lstm = mode == "LSTM"
 
         self.zeros = P.Zeros()

tensorlayerx/nn/core/common.py

@@ -6,6 +6,7 @@
 from tensorlayerx.files import utils
 from tensorlayerx import logging
 import numpy as np
+from queue import Queue
 from tensorlayerx.nn.initializers import *
 
 if tlx.BACKEND == 'mindspore':
@@ -445,4 +446,112 @@ def tolist(tensors):
             ntensors += tolist(t)
         return ntensors
     else:
-        return [tensors]
+        return [tensors]
+
+def construct_graph(inputs, outputs):
+    """construct computation graph for model using ModuleNode object"""
+    all_layers = []
+    node_by_depth = []
+
+    input_tensors_list = inputs if isinstance(inputs, list) else inputs
+
+    queue_node = Queue()
+    # BFS to visit all nodes that should be involved in the computation graph
+    output_tensors_list = outputs if isinstance(outputs, list) else [outputs]
+    output_nodes = [tensor._info[0] for tensor in output_tensors_list]
+
+    visited_node_names = set()
+    for out_node in output_nodes:
+        if out_node.visited:
+            continue
+        queue_node.put(out_node)
+
+        while not queue_node.empty():
+            cur_node = queue_node.get()
+            in_nodes = cur_node.in_nodes
+
+            for node in in_nodes:
+                node.out_nodes.append(cur_node)
+                if not node.visited:
+                    queue_node.put(node)
+                    node.visited = True
+                    if node.node_name not in visited_node_names:
+                        visited_node_names.add(node.node_name)
+                    # else have multiple layers with the same name
+                    else:
+                        raise ValueError(
+                            'Layer name \'%s\' has already been used by another layer. Please change the layer name.'
+                            % node.layer.name
+                        )
+
+    # construct the computation graph in top-sort order
+    cur_depth = [tensor._info[0] for tensor in input_tensors_list]
+    next_depth = []
+    indegrees = {}
+
+    visited_layer_names = []
+    while not len(cur_depth) == 0:
+        node_by_depth.append(cur_depth)
+        for node in cur_depth:
+            if node.layer.name not in visited_layer_names:
+                all_layers.append(node.layer)
+                visited_layer_names.append(node.layer.name)
+            for out_node in node.out_nodes:
+                if out_node.node_name not in indegrees.keys():
+                    indegrees[out_node.node_name] = len(out_node.in_nodes)
+                indegrees[out_node.node_name] -= 1
+                if indegrees[out_node.node_name] == 0:
+                    next_depth.append(out_node)
+        cur_depth = next_depth
+        next_depth = []
+    return node_by_depth, all_layers
+
+
+class ModuleNode(object):
+    """
+    The class :class:`ModuleNode` class represents a conceptional node for a layer.
+
+    ModuleNode is used for building topology and it is actually a light weighted
+    wrapper over Layer.
+
+    Parameters
+    ----------
+    layer : tl.layers.Layer
+        A tl layer that wants to create a node.
+    node_index : int
+        Index of this node in layer._nodes.
+    in_nodes ：a list of ModuleNode
+        Father nodes to this node.
+    in_tensors : a list of tensors
+        Input tensors to this node.
+    out_tensors : a list of tensors
+        Output tensors to this node.
+    in_tensor_idxes : a list of int
+        Indexes of each input tensor in its corresponding node's out_tensors.
+
+    Methods
+    ---------
+    __init__()
+        Initializing the ModuleNode.
+    __call__()
+        (1) Forwarding through the layer. (2) Update its input/output tensors.
+    """
+
+    def __init__(self, layer, node_index, in_nodes, in_tensors, out_tensors, in_tensor_idxes):
+        self.layer = layer
+        self.node_index = node_index
+        self.in_nodes = in_nodes
+        self.out_nodes = []
+        self.in_tensors = in_tensors
+        self.out_tensors = out_tensors
+        self.node_name = layer.name + "_node_{}".format(node_index)
+
+        self.in_tensors_idxes = in_tensor_idxes
+        self.visited = False
+
+    def __call__(self, inputs, **kwargs):
+        """(1) Forwarding through the layer. (2) Update its input/output tensors."""
+        outputs = self.layer(inputs, **kwargs)
+        self.in_tensors = tolist(inputs)
+        self.out_tensors = tolist(outputs)
+        return self.out_tensors

tensorlayerx/nn/core/core_mindspore.py

@@ -1,11 +1,11 @@
 #! /usr/bin/python
 # -*- coding: utf-8 -*-
 
-from .common import str2act, str2init, random_normal
+from .common import str2act, str2init, random_normal, tolist, construct_graph, ModuleNode
 from .common import _save_weights, _load_weights, _save_standard_weights_dict, _load_standard_weights_dict
 from mindspore.nn import Cell
 import tensorlayerx as tlx
-from collections import OrderedDict
+import mindspore as ms
 from mindspore import log as logger
 import inspect
 from mindspore import context
@@ -16,6 +16,7 @@
 __all__ = ['Module', 'Sequential', 'ModuleList', 'ModuleDict']
 
 _global_layer_name_dict = {}
+_global_layer_node = []
 
 
 class Module(Cell):
@@ -138,41 +139,40 @@ def _compute_shape(tensors):
             shape_mem = tlx.get_tensor_shape(tensors)
         return shape_mem
 
-    def __call__(self, *inputs, **kwargs):
+    def __call__(self, *args, **kwargs):
         if self.__class__.construct is Cell.construct:
-            logger.warning(
-                f"The '{self.__class__}' does not override the method 'construct', "
-                f"will call the super class(Cell) 'construct'."
-            )
+            logger.warning(f"The '{self.__class__}' does not override the method 'construct', "
+                           f"will call the super class(Cell) 'construct'.")
         if kwargs:
-            bound_args = inspect.signature(self.construct).bind(*inputs, **kwargs)
-            inputs = bound_args.args
-            kwargs = bound_args.kwargs
+            bound_arguments = inspect.signature(self.construct).bind(*args, **kwargs)
+            bound_arguments.apply_defaults()
+            args = bound_arguments.args
+            kwargs = bound_arguments.kwargs
 
         # Run in Graph mode.
-        if context.get_context("mode") == context.GRAPH_MODE:
-            raise NotImplemented("GRAPH MODE is not supported, please select PYNATIVE MODE.")
+        if context._get_mode() == context.GRAPH_MODE:
+            self._check_construct_args(*args, **kwargs)
+            if self.enable_hook:
+                raise ValueError("For 'Cell', it's not support hook function in graph mode, please use "
+                                 "context.set_context to set pynative mode.")
+            out = self.compile_and_run(*args)
+            return out
 
         # Run in PyNative mode.
         if _pynative_executor.is_top_cell():
             _pynative_executor.set_lazy_build(True)
             # There many Casts in parameter_broadcast. Enable lazy_build and build faster.
             self._do_parameter_broadcast()
 
-        for item in inputs:
-            if isinstance(item, numpy.ndarray):
-                raise TypeError("The cell inputs should not be numpy arrays.")
+        for item in args:
+            if isinstance(item, ms.Tensor) and item.has_init:
+                item.init_data()
+            elif isinstance(item, numpy.ndarray):
+                raise TypeError("For 'Cell', inputs should not be numpy array.")
         if self.requires_grad is True:
             _pynative_executor.set_grad_flag(True)
-        _pynative_executor.new_graph(self, *inputs, **kwargs)
-        cast_inputs = list()
-        if hasattr(self, "_mindspore_flags"):
-            if self._mindspore_flags.get('fp16'):
-                cast_inputs = self._cast_mixed_precision_inputs(inputs, tlx.float16)
-            if self._mindspore_flags.get('fp32'):
-                cast_inputs = self._cast_mixed_precision_inputs(inputs, tlx.float32)
-        if not cast_inputs:
-            cast_inputs = inputs
+        _pynative_executor.new_graph(self, *args, **kwargs)
+        cast_inputs = self.auto_cast_inputs(args)
 
         with self.CellGuard():
             try:
@@ -182,29 +182,13 @@ def __call__(self, *inputs, **kwargs):
                 raise err
 
         if _pynative_executor.is_top_cell():
-            _pynative_executor.execute_all_task()
+            _pynative_executor.execute_lazy_task()
 
         if isinstance(output, Parameter):
             output = output.data
-        _pynative_executor.end_graph(self, output, *inputs, **kwargs)
+        _pynative_executor.end_graph(self, output, *args, **kwargs)
         return output
 
-    def _add_node(self, input_tensors, output_tensors):
-        """Add a LayerNode for this layer given input_tensors, output_tensors.
-
-        WARINING: This function should not be called from outside, it should only be called
-        in layer.__call__ when building static model.
-
-        Parameters
-        ----------
-        input_tensors : Tensor or a list of tensors
-            Input tensors to this layer.
-        output_tensors : Tensor or a list of tensors
-            Output tensors to this layer.
-
-        """
-        raise NotImplementedError
-
     def set_train(self):
         """
         Sets the cell to training mode.
@@ -310,6 +294,49 @@ def init_build(self, *inputs, **kwargs):
 
         self.forward(*inputs, **kwargs)
 
+    def build_graph(self, *inputs, **kwargs):
+        # Add nodes only when the composition is needed.
+        layers = self.cells_and_names(name_prefix='')
+        for layer_name, layer in layers:
+            if isinstance(layer, Module):
+                layer._build_graph = True
+
+        outputs = self.forward(*inputs, **kwargs)
+        self.inputs = inputs
+        self.outputs = outputs
+        self._node_by_depth, self._all_layers = construct_graph(self.inputs, self.outputs)
+        return self._node_by_depth, self._all_layers
+
+    def _add_node(self, input_tensors, output_tensors):
+        """Add a ModuleNode for this layer given input_tensors, output_tensors.
+
+        This function should not be called from outside, it should only be called
+        in __call__ when building static model.
+
+        Parameters
+        ----------
+        input_tensors : Tensor or a list of tensors
+            Input tensors to this layer.
+        output_tensors : Tensor or a list of tensors
+            Output tensors to this layer.
+
+        """
+        inputs_list = tolist(input_tensors)
+        outputs_list = tolist(output_tensors)
+        if self.__class__.__name__ in tlx.layers.inputs.__all__:
+            # for InputLayer, there should be no in_nodes
+            in_nodes = []
+            in_tensor_idxes = [0]
+        else:
+            in_nodes = [tensor._info[0] for tensor in inputs_list]
+            in_tensor_idxes = [tensor._info[1] for tensor in inputs_list]
+        node_index = len(_global_layer_node)
+
+        new_node = ModuleNode(self, node_index, in_nodes, inputs_list, outputs_list, in_tensor_idxes)
+        _global_layer_node.append(new_node)
+        for idx, tensor in enumerate(outputs_list):
+            tensor._info = (new_node, idx)
+
 
 class Sequential(Module):
     """

tensorlayerx/nn/core/core_paddle.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 import copy, six
-from .common import str2act, str2init
+from .common import str2act, str2init, tolist, construct_graph, ModuleNode
 from .common import _save_weights, _load_weights, _save_standard_weights_dict, _load_standard_weights_dict
 from paddle.fluid import framework
 from paddle.fluid.dygraph import Layer
@@ -12,7 +12,11 @@
 import paddle as pd
 from collections import OrderedDict
 from collections import OrderedDict, abc as container_abcs
+import tensorlayerx as tlx
+from queue import Queue
+
 _global_layer_name_dict = {}
+_global_layer_node = []
 
 __all__ = ['Module', 'Sequential', 'ModuleList', 'ModuleDict']
 
@@ -305,6 +309,50 @@ def insert_child_to_layer(self, child_name, child):
             raise TypeError("Child layer type is incorrect.")
         self._sub_layers[child_name] = child
 
+    def build_graph(self, *inputs, **kwargs):
+        # Add nodes only when the composition is needed.
+        # for layer in self.sublayers():
+        #     if isinstance(layer, Module):
+        #         layer._build_graph = True
+        #
+        # outputs = self.forward(*inputs, **kwargs)
+        # self.inputs = inputs
+        # self.outputs = outputs
+        # self._node_by_depth, self._all_layers = construct_graph(self.inputs, self.outputs)
+        # return self._node_by_depth, self._all_layers
+        raise NotImplementedError
+
+    def _add_node(self, input_tensors, output_tensors):
+        """Add a ModuleNode for this layer given input_tensors, output_tensors.
+
+        This function should not be called from outside, it should only be called
+        in __call__ when building static model.
+
+        Parameters
+        ----------
+        input_tensors : Tensor or a list of tensors
+            Input tensors to this layer.
+        output_tensors : Tensor or a list of tensors
+            Output tensors to this layer.
+
+        """
+        pass
+        # inputs_list = tolist(input_tensors)
+        # outputs_list = tolist(output_tensors)
+        # if self.__class__.__name__ in tlx.layers.inputs.__all__:
+        #     # for InputLayer, there should be no in_nodes
+        #     in_nodes = []
+        #     in_tensor_idxes = [0]
+        # else:
+        #     in_nodes = [tensor[0] for tensor in inputs_list]
+        #     in_tensor_idxes = [tensor[1] for tensor in inputs_list]
+        # node_index = len(_global_layer_node)
+        #
+        # new_node = ModuleNode(self, node_index, in_nodes, inputs_list, outputs_list, in_tensor_idxes)
+        # _global_layer_node.append(new_node)
+        # for idx, tensor in enumerate(outputs_list):
+        #     tensor._info = (new_node, idx)
+
 
 class Sequential(Module):