diff --git a/src/models/neural_networks/ealstm.py b/src/models/neural_networks/ealstm.py index a091ec2a4..6f8f0b1f2 100644 --- a/src/models/neural_networks/ealstm.py +++ b/src/models/neural_networks/ealstm.py @@ -254,22 +254,31 @@ def __init__( if dense_features[-1] != 1: dense_features.append(1) - self.dense_layers = nn.ModuleList( - [ + # add linear layer with nonlinear activation functions + dense_layers = [] + for i in range(1, len(dense_features)): + dense_layers.append( nn.Linear( - in_features=dense_features[i - 1], out_features=dense_features[i] + # in = size of previous dense layer + in_features=dense_features[i - 1], + # out = size of current dense layer + out_features=dense_features[i], ) - for i in range(1, len(dense_features)) - ] - ) + ) + if i < len(dense_features) - 1: + # add a ReLU to all layers except the final layer + dense_layers.append(nn.ReLU()) + + self.dense_layers = nn.ModuleList(dense_layers) self.initialize_weights() def initialize_weights(self): - for dense_layer in self.dense_layers: - nn.init.kaiming_uniform_(dense_layer.weight.data) - nn.init.constant_(dense_layer.bias.data, 0) + # initialise weights for all linear layers + if not isinstance(dense_layer, nn.ReLU): + nn.init.kaiming_uniform_(dense_layer.weight.data) + nn.init.constant_(dense_layer.bias.data, 0) def forward( self,