[MXNET-1431] Multiple channel support in Gluon PReLU

python/mxnet/gluon/nn/activations.py

@@ -120,18 +120,23 @@ class PReLU(HybridBlock):
     ----------
     alpha_initializer : Initializer
         Initializer for the `embeddings` matrix.
-
+    in_channels : int, default 1
+        Number of channels (alpha parameters) to learn. Can either be 1
+        or `n` where `n` is the size of the second dimension of the input
+        tensor.
 
     Inputs:
         - **data**: input tensor with arbitrary shape.
 
     Outputs:
         - **out**: output tensor with the same shape as `data`.
     """
-    def __init__(self, alpha_initializer=initializer.Constant(0.25), **kwargs):
+    def __init__(self, alpha_initializer=initializer.Constant(0.25),
+                 in_channels=1, **kwargs):
         super(PReLU, self).__init__(**kwargs)
         with self.name_scope():
-            self.alpha = self.params.get('alpha', shape=(1,), init=alpha_initializer)
+            self.alpha = self.params.get('alpha', shape=(in_channels,),
+                                         init=alpha_initializer)
 
     def hybrid_forward(self, F, x, alpha):
         return F.LeakyReLU(x, gamma=alpha, act_type='prelu', name='fwd')

tests/python/unittest/test_gluon.py

@@ -1240,6 +1240,11 @@ def selu(x):
     x = point_to_validate.reshape((1, 3, 2))
     assert_almost_equal(prelu(x).asnumpy(), mx.nd.where(x >= 0, x, 0.25 * x).asnumpy())
 
+    multichannel_init = mx.initializer.Constant(mx.nd.array([0.1, 0.25, 0.5]))
+    prelu_multichannel = mx.gluon.nn.PReLU(alpha_initializer=multichannel_init, in_channels=3)
+    prelu_multichannel.initialize()
+    assert_almost_equal(prelu_multichannel(x).asnumpy(), np.array([[-0.01, 0.1], [-0.025, 0.1], [-0.05, 0.1]]))
+
     gelu = mx.gluon.nn.GELU()
     def gelu_test(x):
         CUBE_CONSTANT = 0.044715