Fix processing activation

Laicheng0830 · Laicheng0830 · commit 5ff86869acd8 · 2022-05-26T17:57:15.000+08:00
diff --git a/examples/basic_tutorials/tutorial_tensorlayerx_graph.py b/examples/basic_tutorials/tutorial_tensorlayerx_graph.py
@@ -0,0 +1,63 @@
+#! /usr/bin/python
+# -*- coding: utf-8 -*-
+
+import tensorlayerx as tlx
+from tensorlayerx.nn import Module
+from tensorlayerx.nn import Linear, Conv2d, BatchNorm2d, MaxPool2d, Flatten
+
+class CNN(Module):
+
+    def __init__(self):
+        super(CNN, self).__init__()
+        # weights init
+        W_init = tlx.nn.initializers.truncated_normal(stddev=5e-2)
+        W_init2 = tlx.nn.initializers.truncated_normal(stddev=0.04)
+        b_init2 = tlx.nn.initializers.constant(value=0.1)
+
+        self.conv1 = Conv2d(64, (5, 5), (1, 1), padding='SAME', W_init=W_init, b_init=None, name='conv1', in_channels=3, act=tlx.nn.ReLU)
+        self.bn = BatchNorm2d(num_features=64, act=tlx.nn.ReLU)
+        self.maxpool1 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool1')
+
+        self.conv2 = Conv2d(
+            64, (5, 5), (1, 1), padding='SAME', act=tlx.nn.ReLU, W_init=W_init, b_init=None, name='conv2', in_channels=64
+        )
+        self.maxpool2 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool2')
+
+        self.flatten = Flatten(name='flatten')
+        self.linear1 = Linear(384, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear1relu', in_features=2304)
+        self.linear2 = Linear(192, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear2relu', in_features=384)
+        self.linear3 = Linear(10, act=None, W_init=W_init2, name='output1', in_features=192)
+        self.linear4 = Linear(20, act=None, W_init=W_init2, name='output2', in_features=192)
+        self.concat = tlx.nn.Concat(name='concat')
+
+    def forward(self, x):
+        z = self.conv1(x)
+        z = self.bn(z)
+        z = self.maxpool1(z)
+        z = self.conv2(z)
+        z = self.maxpool2(z)
+        z = self.flatten(z)
+        z = self.linear1(z)
+        z = self.linear2(z)
+        z1 = self.linear3(z)
+        z2 = self.linear4(z)
+        z = self.concat([z1, z2])
+        return z
+
+model = CNN()
+inputs = tlx.nn.Input(shape=(3, 24, 24, 3))
+outputs = model(inputs)
+
+node_by_depth, all_layers = model.build_graph(inputs)
+
+for depth, nodes in enumerate(node_by_depth):
+    if depth == 0:
+        if isinstance(inputs, list):
+            assert len(inputs) == len(nodes)
+            for idx, node in enumerate(nodes):
+                print(node.node_name, node.layer)
+        else:
+            print(nodes[0].node_name, nodes[0].layer)
+    else:
+        for node in nodes:
+            print(node.node_name, node.layer)
diff --git a/tensorlayerx/nn/core/common.py b/tensorlayerx/nn/core/common.py
@@ -87,6 +87,28 @@ def str2act(act):
     return _act_dict[act]
 
 
+def processing_act(act):
+    # Processing strings as input, activation functions without parameters。
+    if isinstance(act, str):
+        str_act = str2act(act)
+    if act:
+        # Processing strings as input, activation functions with parameters。
+        if isinstance(act, str) and (len(act) > 5 and act[0:5] == "lrelu" or
+                                     len(act) > 10 and act[0:10] == "leaky_relu"):
+            out_act = str_act
+        elif isinstance(act, str):
+            out_act = str_act()
+        # Processing classes or functions as input, activation functions without parameters
+        elif type(act) == type(tlx.nn.ReLU):
+            out_act = act()
+        # Processing class or function as input, activation function with parameters
+        else:
+            out_act = act
+    else:
+        out_act = act
+    return out_act
+
+
 def _save_weights(net, file_path, format=None):
     """Input file_path, save model weights into a file of given format.
                 Use net.load_weights() to restore.
diff --git a/tensorlayerx/nn/core/core_mindspore.py b/tensorlayerx/nn/core/core_mindspore.py
@@ -1,7 +1,7 @@
 #! /usr/bin/python
 # -*- coding: utf-8 -*-
 
-from .common import check_parameter, str2act, str2init, random_normal, tolist, construct_graph, ModuleNode, select_attrs
+from .common import check_parameter, processing_act, str2init, random_normal, tolist, construct_graph, ModuleNode, select_attrs
 from .common import _save_weights, _load_weights, _save_standard_weights_dict, _load_standard_weights_dict
 from mindspore.nn import Cell
 import tensorlayerx as tlx
@@ -49,19 +49,7 @@ def __init__(self, name=None, act=None, *args, **kwargs):
 
         self.name = name
 
-        if isinstance(act, str):
-            str_act = str2act(act)
-
-        if act:
-            if isinstance(act, str) and (len(act) > 5 and act[0:5] == "lrelu" or
-                                         len(act) > 10 and act[0:10] == "leaky_relu"):
-                self.act = str_act
-            elif isinstance(act, str):
-                self.act = str_act()
-            else:
-                self.act = act()
-        else:
-            self.act = act
+        self.act = processing_act(act)
 
         # Layer building state
         self._built = False
diff --git a/tensorlayerx/nn/core/core_paddle.py b/tensorlayerx/nn/core/core_paddle.py
@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 import copy, six
-from .common import check_parameter, str2act, str2init, tolist, construct_graph, ModuleNode
+from .common import check_parameter, processing_act, str2init
 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
@@ -49,19 +49,7 @@ def __init__(self, name=None, act=None, *args, **kwargs):
 
         self.name = name
 
-        if isinstance(act, str):
-            str_act = str2act(act)
-
-        if act:
-            if isinstance(act, str) and (len(act) > 5 and act[0:5] == "lrelu" or
-                                         len(act) > 10 and act[0:10] == "leaky_relu"):
-                self.act = str_act
-            elif isinstance(act, str):
-                self.act = str_act()
-            else:
-                self.act = act()
-        else:
-            self.act = act
+        self.act = processing_act(act)
 
         # Layer building state
         self._built = False
diff --git a/tensorlayerx/nn/core/core_tensorflow.py b/tensorlayerx/nn/core/core_tensorflow.py
@@ -1,7 +1,7 @@
 #! /usr/bin/python
 # -*- coding: utf-8 -*-
 
-from .common import check_parameter, str2act, str2init, tolist, construct_graph, ModuleNode, select_attrs
+from .common import check_parameter, processing_act, str2init, tolist, construct_graph, ModuleNode, select_attrs
 from .common import _save_weights, _load_weights, _save_standard_weights_dict, _load_standard_weights_dict
 from collections import OrderedDict, abc as container_abcs
 from collections import OrderedDict
@@ -82,19 +82,7 @@ def __init__(self, name=None, act=None, *args, **kwargs):
 
         self.name = name
 
-        if isinstance(act, str):
-            str_act = str2act(act)
-
-        if act:
-            if isinstance(act, str) and (len(act) > 5 and act[0:5] == "lrelu" or
-                                         len(act) > 10 and act[0:10] == "leaky_relu"):
-                self.act = str_act
-            elif isinstance(act, str):
-                self.act = str_act()
-            else:
-                self.act = act()
-        else:
-            self.act = act
+        self.act = processing_act(act)
 
         # Layer building state
         self._built = False
@@ -587,8 +575,7 @@ def check_param(self, param, dim='2d'):
 
     def build_graph(self, *inputs, **kwargs):
         # Add nodes only when the composition is needed.
-        layers = self.layers_and_names(name_prefix='')
-        for layer_name, layer in layers:
+        for layer_name, layer in self._layers.items():
             if isinstance(layer, Module):
                 layer._build_graph = True
         self.set_eval()
diff --git a/tensorlayerx/nn/core/core_torch.py b/tensorlayerx/nn/core/core_torch.py
@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 from torch.nn import Module as T_Module
-from .common import check_parameter, str2act, str2init, tolist, construct_graph, ModuleNode, select_attrs
+from .common import check_parameter, processing_act, str2init, tolist, construct_graph, ModuleNode, select_attrs
 from .common import _save_weights, _load_weights, _save_standard_weights_dict, _load_standard_weights_dict
 from torch.nn.parameter import Parameter
 from collections import OrderedDict
@@ -46,19 +46,7 @@ def __init__(self, name=None, act=None, *args, **kwargs):
 
         self.name = name
 
-        if isinstance(act, str):
-            str_act = str2act(act)
-
-        if act:
-            if isinstance(act, str) and (len(act) > 5 and act[0:5] == "lrelu" or
-                                         len(act) > 10 and act[0:10] == "leaky_relu"):
-                self.act = str_act
-            elif isinstance(act, str):
-                self.act = str_act()
-            else:
-                self.act = act()
-        else:
-            self.act = act
+        self.act = processing_act(act)
 
         # Layer building state
         self._built = False
