Merge pull request #1038 from tensorlayer/Jingqing-patch

minor update lambda layer doc and test

Author: JingqingZ

tensorlayer/layers/lambda_layers.py

@@ -35,14 +35,14 @@ class Lambda(Layer):
 
     Examples
     ---------
-    Non-parametric and non-args case
+    Non-parametric and non-args case:
     This case is supported in the Model.save() / Model.load() to save / load the whole model architecture and weights(optional).
 
     >>> x = tl.layers.Input([8, 3], name='input')
     >>> y = tl.layers.Lambda(lambda x: 2*x, name='lambda')(x)
 
 
-    Non-parametric and with args case
+    Non-parametric and with args case:
     This case is supported in the Model.save() / Model.load() to save / load the whole model architecture and weights(optional).
 
     >>> def customize_func(x, foo=42): # x is the inputs, foo is an argument
@@ -51,19 +51,19 @@ class Lambda(Layer):
     >>> lambdalayer = tl.layers.Lambda(customize_func, fn_args={'foo': 2}, name='lambda')(x)
 
 
-    Any function with outside variables
+    Any function with outside variables:
     This case has not been supported in Model.save() / Model.load() yet.
     Please avoid using Model.save() / Model.load() to save / load models that contain such Lambda layer. Instead, you may use Model.save_weights() / Model.load_weights() to save / load model weights.
     Note: In this case, fn_weights should be a list, and then the trainable weights in this Lambda layer can be added into the weights of the whole model.
 
-    >>> vara = [tf.Variable(1.0)]
+    >>> a = tf.Variable(1.0)
     >>> def func(x):
-    >>>     return x + vara
+    >>>     return x + a
     >>> x = tl.layers.Input([8, 3], name='input')
-    >>> y = tl.layers.Lambda(func, fn_weights=a, name='lambda')(x)
+    >>> y = tl.layers.Lambda(func, fn_weights=[a], name='lambda')(x)
 
 
-    Parametric case, merge other wrappers into TensorLayer
+    Parametric case, merge other wrappers into TensorLayer:
     This case is supported in the Model.save() / Model.load() to save / load the whole model architecture and weights(optional).
 
     >>> layers = [
@@ -74,27 +74,27 @@ class Lambda(Layer):
     >>> perceptron = tf.keras.Sequential(layers)
     >>> # in order to compile keras model and get trainable_variables of the keras model
     >>> _ = perceptron(np.random.random([100, 5]).astype(np.float32))
-
+    >>>
     >>> class CustomizeModel(tl.models.Model):
     >>>     def __init__(self):
     >>>         super(CustomizeModel, self).__init__()
     >>>         self.dense = tl.layers.Dense(in_channels=1, n_units=5)
     >>>         self.lambdalayer = tl.layers.Lambda(perceptron, perceptron.trainable_variables)
-
+    >>>
     >>>     def forward(self, x):
     >>>         z = self.dense(x)
     >>>         z = self.lambdalayer(z)
     >>>         return z
-
+    >>>
     >>> optimizer = tf.optimizers.Adam(learning_rate=0.1)
     >>> model = CustomizeModel()
     >>> model.train()
-
+    >>>
     >>> for epoch in range(50):
     >>>     with tf.GradientTape() as tape:
     >>>         pred_y = model(data_x)
     >>>         loss = tl.cost.mean_squared_error(pred_y, data_y)
-
+    >>>
     >>>     gradients = tape.gradient(loss, model.trainable_weights)
     >>>     optimizer.apply_gradients(zip(gradients, model.trainable_weights))
 

tests/layers/test_layers_lambda.py

@@ -101,6 +101,32 @@ def forward(self, x, bar):
         out, out2 = model(self.data_x, bar=2)
         self.assertTrue(np.array_equal(out2.numpy(), out.numpy()))
 
+    def test_lambda_func_with_weight(self):
+
+        a = tf.Variable(1.0)
+
+        def customize_fn(x):
+            return x + a
+
+        class CustomizeModel(tl.models.Model):
+
+            def __init__(self):
+                super(CustomizeModel, self).__init__()
+                self.dense = tl.layers.Dense(in_channels=1, n_units=5)
+                self.lambdalayer = tl.layers.Lambda(customize_fn, fn_weights=[a])
+
+            def forward(self, x):
+                z = self.dense(x)
+                z = self.lambdalayer(z)
+                return z
+
+        model = CustomizeModel()
+        print(model.lambdalayer)
+        model.train()
+
+        out = model(self.data_x)
+        print(out.shape)
+
     def test_lambda_func_without_args(self):
 
         class CustomizeModel(tl.models.Model):