Update weight and bias for selu activation

vulcanai2/models/layers.py

@@ -2,6 +2,7 @@
 import torch
 import torch.nn as nn
 import logging
+import numpy as np
 logger = logging.getLogger(__name__)
 
 
@@ -52,8 +53,24 @@ def _init_weights(self):
         if self.initializer is None, then pytorch default weight
         will be assigned to the kernel
         """
+        #import pdb; pdb.set_trace()
         if self.initializer:
             self.initializer(self._kernel.weight)
+        elif isinstance(self.activation, nn.SELU):
+            stdv = np.sqrt(1. / self.in_features)
+            weight = torch.empty(self.out_features, self.in_features)
+            weight = nn.init.normal_(weight, stdv)
+            self._kernel.weight = nn.Parameter(weight)
+        elif isinstance(self.activation, nn.ReLU):
+            weight = torch.empty(self.out_features, self.in_features)
+            weight = nn.init.kaiming_normal_(weight, mode='fan_in', nonlinearity=self.activation)
+            self._kernel.weight = nn.Parameter(weight)
+        else:
+            weight = torch.empty(self.out_features, self.in_features)
+            weight = nn.init.xavier_uniform_(weight)
+
+            self._kernel.weight = nn.Parameter(weight)
+        #import pdb; pdb.set_trace()
 
     def _init_bias(self):
         """
@@ -64,7 +81,14 @@ def _init_bias(self):
         """
         if self.bias_init:
             nn.init.constant_(self._kernel.bias, self.bias_init)
-
+        elif isinstance(self.activation, nn.SELU):
+            bias = torch.empty(self.out_features)
+            bias = nn.init.normal_(bias, 0.0)
+            self._kernel.bias = nn.Parameter(bias)
+        else:
+            bias = torch.empty(self.out_features)
+            bias = nn.init.constant_(bias, 0.0)
+            self._kernel.bias = nn.Parameter(bias)
 
 class FlattenUnit(BaseUnit):
     """
@@ -121,17 +145,16 @@ def __init__(self, in_features, out_features,
                                         norm, dropout)
         self.in_features = in_features
         self.out_features = out_features
-
+        self.activation = activation
         # Main layer
         self._kernel = nn.Linear(
                             in_features=self.in_features,
                             out_features=self.out_features,
                             bias=True
                             )
         self.add_module('_kernel', self._kernel)
-        self._init_weights()
-        self._init_bias()
-
+        #self._init_weights()
+        #self._init_bias()
         # Norm
         if self.norm is not None:
             if self.norm == 'batch':
@@ -155,7 +178,8 @@ def __init__(self, in_features, out_features,
             else:
                 self.add_module(
                     '_dropout', nn.Dropout(self.dropout))
-
+        self._init_weights()
+        self._init_bias()
 
 # TODO: Automatically calculate padding to be the same as input shape.
 class ConvUnit(BaseUnit):