PELU fixes and upgrade

keras_contrib/layers/advanced_activations.py

@@ -1,5 +1,8 @@
 from .. import initializers
+from .. import regularizers
+from .. import constraints
 from keras.engine import Layer
+from keras.engine import InputSpec
 from .. import backend as K
 from keras.utils.generic_utils import get_custom_objects
 import numpy as np
@@ -33,51 +36,79 @@ class PELU(Layer):
         - [PARAMETRIC EXPONENTIAL LINEAR UNIT FOR DEEP CONVOLUTIONAL NEURAL NETWORKS](https://arxiv.org/abs/1605.09332v3)
     """
 
-    def __init__(self, alphas_initializer='one', betas_initializer='one', weights=None, shared_axes=None, **kwargs):
+    def __init__(self, alpha_initializer='ones',
+                 alpha_regularizer=None,
+                 alpha_constraint=None,
+                 beta_initializer='ones',
+                 beta_regularizer=None,
+                 beta_constraint=None,
+                 shared_axes=None,
+                 **kwargs):
+        super(PELU, self).__init__(**kwargs)
         self.supports_masking = True
-        self.alphas_initializer = initializers.get(alphas_initializer)
-        self.betas_initializer = initializers.get(betas_initializer)
-        self.initial_weights = weights
-        if not isinstance(shared_axes, (list, tuple)):
+        self.alpha_initializer = initializers.get(alpha_initializer)
+        self.alpha_regularizer = regularizers.get(alpha_regularizer)
+        self.alpha_constraint = constraints.get(alpha_constraint)
+        self.beta_initializer = initializers.get(beta_initializer)
+        self.beta_regularizer = regularizers.get(beta_regularizer)
+        self.beta_constraint = constraints.get(beta_constraint)
+        if shared_axes is None:
+            self.shared_axes = None
+        elif not isinstance(shared_axes, (list, tuple)):
             self.shared_axes = [shared_axes]
         else:
             self.shared_axes = list(shared_axes)
-        super(PELU, self).__init__(**kwargs)
 
     def build(self, input_shape):
         param_shape = list(input_shape[1:])
         self.param_broadcast = [False] * len(param_shape)
-        if self.shared_axes[0] is not None:
+        if self.shared_axes is not None:
             for i in self.shared_axes:
                 param_shape[i - 1] = 1
                 self.param_broadcast[i - 1] = True
 
         # Initialised as ones to emulate the default ELU
-        self.alphas = self.add_weight(param_shape,
-                                      name='alpha',
-                                      initializer=self.alphas_initializer)
-        self.betas = self.add_weight(param_shape, name='betas', initializer=self.betas_initializer)
+        self.alpha = self.add_weight(param_shape,
+                                     name='alpha',
+                                     initializer=self.alpha_initializer,
+                                     regularizer=self.alpha_regularizer,
+                                     constraint=self.alpha_constraint)
+        self.beta = self.add_weight(param_shape,
+                                    name='beta',
+                                    initializer=self.beta_initializer,
+                                    regularizer=self.beta_regularizer,
+                                    constraint=self.beta_constraint)
 
-        self.trainable_weights = [self.alphas, self.betas]
-
-        if self.initial_weights is not None:
-            self.set_weights(self.initial_weights)
-            del self.initial_weights
+        # Set input spec
+        axes = {}
+        if self.shared_axes:
+            for i in range(1, len(input_shape)):
+                if i not in self.shared_axes:
+                    axes[i] = input_shape[i]
+        self.input_spec = InputSpec(ndim=len(input_shape), axes=axes)
+        self.built = True
 
     def call(self, x, mask=None):
         if K.backend() == 'theano':
-            pos = K.relu(x) * (K.pattern_broadcast(self.alphas, self.param_broadcast) /
-                               K.pattern_broadcast(self.betas, self.param_broadcast))
-            neg = (K.pattern_broadcast(self.alphas, self.param_broadcast) *
-                   (K.exp((-K.relu(-x)) / K.pattern_broadcast(self.betas, self.param_broadcast)) - 1))
+            pos = K.relu(x) * (K.pattern_broadcast(self.alpha, self.param_broadcast) /
+                               K.pattern_broadcast(self.beta, self.param_broadcast))
+            neg = (K.pattern_broadcast(self.alpha, self.param_broadcast) *
+                   (K.exp((-K.relu(-x)) / K.pattern_broadcast(self.beta, self.param_broadcast)) - 1))
         else:
-            pos = K.relu(x) * self.alphas / self.betas
-            neg = self.alphas * (K.exp((-K.relu(-x)) / self.betas) - 1)
+            pos = K.relu(x) * self.alpha / self.beta
+            neg = self.alpha * (K.exp((-K.relu(-x)) / self.beta) - 1)
         return neg + pos
 
     def get_config(self):
-        config = {'alphas_initializer': initializers.serialize(self.alphas_initializer),
-                  'betas_initializer': initializers.serialize(self.betas_initializer)}
+        config = {
+            'alpha_initializer': initializers.serialize(self.alpha_initializer),
+            'alpha_regularizer': regularizers.serialize(self.alpha_regularizer),
+            'alpha_constraint': constraints.serialize(self.alpha_constraint),
+            'beta_initializer': initializers.serialize(self.beta_initializer),
+            'beta_regularizer': regularizers.serialize(self.beta_regularizer),
+            'beta_constraint': constraints.serialize(self.beta_constraint),
+            'shared_axes': self.shared_axes
+        }
         base_config = super(PELU, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))