fix TLX nodes collection bug

Author: hanjr92

tensorlayerx/nn/core/core_mindspore.py

@@ -241,11 +241,16 @@ def init_build(self, *inputs, **kwargs):
 
         self.forward(*inputs, **kwargs)
 
-    def build_graph(self, *inputs, **kwargs):
-        # Add nodes only when the composition is needed.
+    def set_build_graph(self):
         for layer_name, layer in self._cells.items():
             if isinstance(layer, Module):
+                if len(layer._cells) > 1:
+                    layer.set_build_graph()
                 layer._build_graph = True
+
+    def build_graph(self, *inputs, **kwargs):
+        # Add nodes only when the composition is needed.
+        self.set_build_graph()
         self.set_eval()
 
         outputs = self.forward(*inputs, **kwargs)
@@ -382,31 +387,26 @@ def build(self, inputs_shape):
 
     def forward(self, input_data):
         for layer in self.layer_list:
-            inputs = input_data
             input_data = layer(input_data)
-            outputs = input_data
-            if not self._nodes_fixed and self._build_graph:
-                self._add_seq_node(inputs, outputs, layer)
-        self._nodes_fixed = True
         return input_data
 
-    def _add_seq_node(self, input_tensors, output_tensors, layer):
-        inputs_list = tolist(input_tensors)
-        outputs_list = tolist(output_tensors)
-        if layer.__class__.__name__ in tlx.layers.inputs.__all__:
-            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(
-            layer, node_index, in_nodes, inputs_list, outputs_list, in_tensor_idxes, select_attrs(layer)
-        )
-        _global_layer_node.append(new_node)
-        for idx, tensor in enumerate(outputs_list):
-            tensor._info = (new_node, idx)
+    # def _add_seq_node(self, input_tensors, output_tensors, layer):
+    #     inputs_list = tolist(input_tensors)
+    #     outputs_list = tolist(output_tensors)
+    #     if layer.__class__.__name__ in tlx.layers.inputs.__all__:
+    #         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(
+    #         layer, node_index, in_nodes, inputs_list, outputs_list, in_tensor_idxes, select_attrs(layer)
+    #     )
+    #     _global_layer_node.append(new_node)
+    #     for idx, tensor in enumerate(outputs_list):
+    #         tensor._info = (new_node, idx)
 
 
 class ModuleList(Module):

tensorlayerx/nn/core/core_paddle.py

@@ -300,6 +300,10 @@ def insert_child_to_layer(self, child_name, child):
             raise TypeError("Child layer type is incorrect.")
         self._sub_layers[child_name] = child
 
+    def set_build_graph(self):
+
+        raise NotImplementedError
+
     def build_graph(self, *inputs, **kwargs):
         # Add nodes only when the composition is needed.
         # for layer in self.sublayers():

tensorlayerx/nn/core/core_tensorflow.py

@@ -639,11 +639,16 @@ def str_to_init(self, initializer):
     def check_param(self, param, dim='2d'):
         return check_parameter(param, dim)
 
-    def build_graph(self, *inputs, **kwargs):
-        # Add nodes only when the composition is needed.
+    def set_build_graph(self):
         for layer_name, layer in self._layers.items():
             if isinstance(layer, Module):
+                if len(layer._layers) > 1:
+                    layer.set_build_graph()
                 layer._build_graph = True
+
+    def build_graph(self, *inputs, **kwargs):
+        # Add nodes only when the composition is needed.
+        self.set_build_graph()
         self.set_eval()
 
         outputs = self.forward(*inputs, **kwargs)
@@ -774,31 +779,26 @@ def build(self, inputs_shape):
 
     def forward(self, input_data):
         for layer in self.layer_list:
-            inputs = input_data
             input_data = layer(input_data)
-            outputs = input_data
-            if not self._nodes_fixed and self._build_graph:
-                self._add_seq_node(inputs, outputs, layer)
-        self._nodes_fixed = True
         return input_data
-
-    def _add_seq_node(self, input_tensors, output_tensors, layer):
-        inputs_list = tolist(input_tensors)
-        outputs_list = tolist(output_tensors)
-        if layer.__class__.__name__ in tlx.layers.inputs.__all__:
-            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(
-            layer, node_index, in_nodes, inputs_list, outputs_list, in_tensor_idxes, select_attrs(layer)
-        )
-        _global_layer_node.append(new_node)
-        for idx, tensor in enumerate(outputs_list):
-            tensor._info = (new_node, idx)
+    #
+    # def _add_seq_node(self, input_tensors, output_tensors, layer):
+    #     inputs_list = tolist(input_tensors)
+    #     outputs_list = tolist(output_tensors)
+    #     if layer.__class__.__name__ in tlx.layers.inputs.__all__:
+    #         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(
+    #         layer, node_index, in_nodes, inputs_list, outputs_list, in_tensor_idxes, select_attrs(layer)
+    #     )
+    #     _global_layer_node.append(new_node)
+    #     for idx, tensor in enumerate(outputs_list):
+    #         tensor._info = (new_node, idx)
 
 
 class ModuleList(Module):

tensorlayerx/nn/core/core_torch.py

@@ -207,11 +207,16 @@ def init_build(self, *inputs, **kwargs):
 
         self.forward(*inputs, **kwargs)
 
-    def build_graph(self, *inputs, **kwargs):
-        # Add nodes only when the composition is needed.
-        for name, layer in self._modules.items():
+    def set_build_graph(self):
+        for layer_name, layer in self._modules.items():
             if isinstance(layer, Module):
+                if len(layer._modules) > 1:
+                    layer.set_build_graph()
                 layer._build_graph = True
+
+    def build_graph(self, *inputs, **kwargs):
+        # Add nodes only when the composition is needed.
+        self.set_build_graph()
         self.set_eval()
 
         outputs = self.forward(*inputs, **kwargs)
@@ -360,31 +365,26 @@ def build(self, inputs_shape):
 
     def forward(self, input_data):
         for layer in self.layer_list:
-            inputs = input_data
             input_data = layer(input_data)
-            outputs = input_data
-            if not self._nodes_fixed and self._build_graph:
-                self._add_seq_node(inputs, outputs, layer)
-        self._nodes_fixed = True
         return input_data
 
-    def _add_seq_node(self, input_tensors, output_tensors, layer):
-        inputs_list = tolist(input_tensors)
-        outputs_list = tolist(output_tensors)
-        if layer.__class__.__name__ in tlx.layers.inputs.__all__:
-            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(
-            layer, node_index, in_nodes, inputs_list, outputs_list, in_tensor_idxes, select_attrs(layer)
-        )
-        _global_layer_node.append(new_node)
-        for idx, tensor in enumerate(outputs_list):
-            tensor._info = (new_node, idx)
+    # def _add_seq_node(self, input_tensors, output_tensors, layer):
+    #     inputs_list = tolist(input_tensors)
+    #     outputs_list = tolist(output_tensors)
+    #     if layer.__class__.__name__ in tlx.layers.inputs.__all__:
+    #         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(
+    #         layer, node_index, in_nodes, inputs_list, outputs_list, in_tensor_idxes, select_attrs(layer)
+    #     )
+    #     _global_layer_node.append(new_node)
+    #     for idx, tensor in enumerate(outputs_list):
+    #         tensor._info = (new_node, idx)
 
 
 class ModuleList(Module):