Made exclusion of activation layer to output function possible, handling class/instances now correctly

Author: RandomDefaultUser

mala/common/parameters.py

@@ -344,26 +344,33 @@ class ParametersNetwork(ParametersBase):
         network. Please note that the input layer is included therein.
         Default: [10,10,0]
 
-    layer_activations : list or str
-        A list of strings detailing the activation functions to be used
-        by the neural network. If a single string is supplied, then this
-        activation function is used for all layers (including the output layer,
-        i.e., an output activation is used!). Otherwise, the activation
-        functions are added layer by layer.
+    layer_activations : list or str or class or nn.Module
+        Detailing the activation functions to be used
+        by the neural network. If a single object is supplied, then this
+        activation function is used for all layers (whether this applies to the
+        output layer is controlled by layer_activations_include_output_layer).
+        Otherwise, the activation functions are added layer by layer.
         Note that no activation function is applied between input layer and
         first hidden layer!
         The items in the list can either be strings, which MALA
-        will map them directly to the correct activation functions or
-        torch.nn.Module objects containing the activation functions directly
-        OR None, in which case no activation function is used. The None
-        can be ommitted at the end, but is useful when layers without
+        will map to the correct activation functions,
+        torch.nn.Module objects, torch.nn.Module classes (which MALA will
+        instantiate) OR None, in which case no activation function is used.
+        The None can be ommitted at the end, but is useful when layers without
         activation functions are to be added in the middle
         Currently supported activation function strings are:
 
             - "Sigmoid"
             - "ReLU"
             - "LeakyReLU" (default)
 
+    layer_activations_include_output_layer : bool
+        If False, no activation function is added to the output layer. This
+        can of course also be done by supplying just the right amount of
+        activation functions and this parameter mainly exist to control the
+        last layer of activation functions in the case of using
+        layer_activations with only a single object.
+
     loss_function_type : string
         Loss function for the neural network
         Currently supported loss functions include:
@@ -398,6 +405,7 @@ def __init__(self):
         self.nn_type = "feed-forward"
         self.layer_sizes = [10, 10, 10]
         self.layer_activations = "LeakyReLU"
+        self.layer_activations_include_output_layer = True
         self.loss_function_type = "mse"
 
         # for LSTM/Gru

mala/network/network.py

@@ -231,15 +231,21 @@ def __init__(self, params):
         # We should NOT modify the list itself. This would break the
         # hyperparameter algorithms.
         use_only_one_activation_type = False
-
-        if not isinstance(self.params.layer_activations, str):
+        if isinstance(self.params.layer_activations, list):
             if len(self.params.layer_activations) > self.number_of_layers:
+
+                number_of_ignored_layers = (
+                    len(self.params.layer_activations) - self.number_of_layers
+                )
+                number_of_ignored_layers += (
+                    1
+                    if self.params.layer_activations_include_output_layer
+                    is False
+                    else 0
+                )
                 printout(
                     "Too many activation layers provided. The last",
-                    str(
-                        len(self.params.layer_activations)
-                        - self.number_of_layers
-                    ),
+                    str(number_of_ignored_layers),
                     "activation function(s) will be ignored.",
                     min_verbosity=1,
                 )
@@ -256,20 +262,24 @@ def __init__(self, params):
                     )
                 )
             )
-            try:
-                if isinstance(self.params.layer_activations, str):
-                    self._append_activation_function(
-                        self.params.layer_activations
-                    )
-                else:
-                    self._append_activation_function(
-                        self.params.layer_activations[i]
-                    )
-            except KeyError:
-                raise Exception("Invalid activation type seleceted.")
-            except IndexError:
-                # No activation functions left to append at the end.
-                pass
+            if (
+                i < self.number_of_layers - 1
+            ) or self.params.layer_activations_include_output_layer:
+                try:
+                    if isinstance(self.params.layer_activations, list):
+                        self._append_activation_function(
+                            self.params.layer_activations[i]
+                        )
+                    else:
+                        self._append_activation_function(
+                            self.params.layer_activations
+                        )
+
+                except KeyError:
+                    raise Exception("Invalid activation type seleceted.")
+                except IndexError:
+                    # No activation functions left to append at the end.
+                    pass
 
         # Once everything is done, we can move the Network on the target
         # device.
@@ -300,7 +310,7 @@ def _append_activation_function(self, activation_function):
 
         Parameters
         ----------
-        activation_function : str
+        activation_function : str or nn.Module or class
             Activation function to be appended.
         """
         if activation_function is None:
@@ -311,6 +321,8 @@ def _append_activation_function(self, activation_function):
             )
         elif isinstance(activation_function, nn.Module):
             self.layers.append(activation_function)
+        elif issubclass(activation_function, nn.Module):
+            self.layers.append(activation_function())
 
 
 class LSTM(Network):