Removed mappings dictionary, now handled over getattr

also updated docs

Author: RandomDefaultUser

mala/common/parameters.py

@@ -352,18 +352,18 @@ class ParametersNetwork(ParametersBase):
         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 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 items in the list can either be strings (=names of torch.nn.Module
+        activation functions), which MALA 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)
-            - "Tanh"
+        activation functions are to be skipped in the middle.
+        Note that output from the output layer is by default restricted to
+        only have positive values via restrict_targets in the ParameterTargets
+        subclass. This is similar to having a ReLU function as a final
+        activation function and ensures the physicality of the outputs (since
+        the (L)DOS can never be negative).
 
     layer_activations_include_output_layer : bool
         If False, no activation function is added to the output layer. This

mala/network/network.py

@@ -94,14 +94,6 @@ def __init__(self, params: Parameters):
         # initialize the parent class
         super(Network, self).__init__()
 
-        # Mappings for parsing of the activation layers.
-        self._activation_mappings = {
-            "Sigmoid": nn.Sigmoid,
-            "ReLU": nn.ReLU,
-            "LeakyReLU": nn.LeakyReLU,
-            "Tanh": nn.Tanh,
-        }
-
         # initialize the layers
         self.number_of_layers = len(self.params.layer_sizes) - 1
 
@@ -316,9 +308,13 @@ def _append_activation_function(self, activation_function):
         if activation_function is None:
             pass
         elif isinstance(activation_function, str):
-            self.layers.append(
-                self._activation_mappings[activation_function]()
-            )
+            try:
+                self.layers.append(getattr(torch.nn, activation_function)())
+            except AttributeError:
+                raise Exception(
+                    "Torch does not contain the specified "
+                    "activation function: " + activation_function
+                )
         elif isinstance(activation_function, nn.Module):
             self.layers.append(activation_function)
         elif issubclass(activation_function, nn.Module):
@@ -367,9 +363,7 @@ def __init__(self, params):
                 self.params.num_hidden_layers,
                 batch_first=True,
             )
-        self.activation = self._activation_mappings[
-            self.params.layer_activations[0]
-        ]()
+        self.activation = getattr(torch.nn, self.params.layer_activations[0])()
 
         self.batch_size = None
         # Once everything is done, we can move the Network on the target
@@ -505,9 +499,7 @@ def __init__(self, params):
                 self.params.num_hidden_layers,
                 batch_first=True,
             )
-        self.activation = self._activation_mappings[
-            self.params.layer_activations[0]
-        ]()
+        self.activation = getattr(torch.nn, self.params.layer_activations[0])()
 
         if params.use_gpu:
             self.to("cuda")