Adding initialization and activation function as parameters in a class

luiscssz: <p>Hi all ,<br>
I am new to Pytorch and need some help.</p>
<p>I would like to add, in the definition of a very simple fully connected NN class (FCN) using only nn.Linear layers:</p>
<ol>
<li>an option to select an activation function (e.g. Tanh, RELU,…) and</li>
<li>a initialization type (Xavier, Kaiming, zeros,…).</li>
</ol>
<p>Here my first code snippet, which unfortunately not works:</p>
<pre><code class="lang-auto">class FCN(nn.Module):   # inherent from nn.Module
    "Defines a standard fully-connected network in PyTorch"
    # Initialization of the neural network layers in __init__:
    def __init__(self, input_dim, hidden_dim, output_dim, activation = "Tanh", initialization = "Xavier_normal"):
        super().__init__()       

        # HERE THE DEFINITION OF ACTIVATION FUNCTIONS:
        activation_functions = nn.ModuleDict([
                                            ["Tanh", nn.Tanh()],
                                            ["ReLU", nn.ReLU()],
                                            ["LeakyReLU", nn.LeakyReLU()],
                                            ["Sigmoid", nn.Sigmoid()],
                                            ["Softmax", nn.Softmax()],
                                            ])
        # HERE THE DEFINITION OF INITIALIZATION OPTION:
        initialization_type = nn.ModuleDict([
                                          ["Xavier_normal", nn.init.xavier_normal_],
                                          ])
        self.initialization = initialization
        
        initialization.get(initialization, None)

        # INPUT Linear function: 
        self.fc_in = nn.Linear(input_dim, hidden_dim)
        
        # INPUT Non-linearity (activation function): 
        self.tangh = nn.Tanh()
        
        # OUTPUT Linear function: 
        self.fc_out = nn.Linear(hidden_dim, output_dim)

        
    def forward(self, input):
        # INPUT Linear function:
        hidden = self.fc_in(input)

        # HIDDEN Non-linearity
        hidden = self.tangh(hidden)

        # OUTPUT Linear function
        output = self.fc_out(hidden)
        
        return output

    def weights_and_bias_initialization(self):
        for module_with_layers in self.modules():
            if isinstance(module_with_layers, nn.Linear):
                self.initialization(module_with_layers.weight)
                #nn.init.xavier_uniform_(module_with_layers.weight)
                #nn.init.xavier_uniform_(module_with_layers.bias)
                #nn.init.constant_(module_with_layers.bias, 0)
</code></pre>
<p>I am using a nn.ModuleDic, following these <a href="https://github.com/FrancescoSaverioZuppichini/Pytorch-how-and-when-to-use-Module-Sequential-ModuleList-and-ModuleDict" rel="noopener nofollow ugc">instructions</a>, by the definition of the posible activation functions (Tanh,…) and <strong>this part works</strong>!!.<br>
The <strong>error</strong> appears by the definition of the initialization using:</p>
<pre><code class="lang-auto">model_test = FCN(2, 8, 1, activation="Tanh", initialization="Xavier_normal")
</code></pre>
<p>when I use again a nn.ModuleDict by the definition of the initialization:</p>
<pre><code class="lang-auto"> initialization_type = nn.ModuleDict([
                                     ["Xavier_normal", nn.init.xavier_normal_],
                                          ])
</code></pre>
<p>ERROR:</p>
<pre><code class="lang-auto">TypeError: torch.nn.init.xavier_normal_ is not a Module subclass
</code></pre>
<p>If I use an instance of the  nn.init.xavier_normal_, I obtain a new error:</p>
<pre><code class="lang-auto">TypeError: xavier_normal_() missing 1 required positional argument: 'tensor'
</code></pre>
<p>I am not quite sure if a nn.ModuleDict is the best option for this case.</p>
<p><strong>A second question related with:</strong></p>
<p><strong>Is the place of the initialization, within the class definition the best place?</strong></p>
<p>I can remember that <a class="mention" href="/u/ptrblck">@ptrblck</a> explained somewhere (unfortunately I can not find again his explanation and a code snippet), that it is better to locate the definition of the function outside of the class after the instance.<br>
If I use the following code snippet, using the success activation option <strong>within</strong> the class and using the activation <strong>outside</strong>:</p>
<pre><code class="lang-auto">def weights_and_bias_initialization(model_instance):
        for module_with_layers in model_instance.modules():
            if isinstance(module_with_layers, nn.Linear):
                nn.init.xavier_uniform_(module_with_layers.weight)
                nn.init.xavier_uniform_(module_with_layers.bias)
class FCN_2(nn.Module):   # inherent from nn.Module
    "Defines a standard fully-connected network in PyTorch"
    # Initialization of the neural network layers in __init__:
    def __init__(self, input_dim, hidden_dim, output_dim, activation = "Tanh"):
        super().__init__()       

        # HERE THE DEFINITION OF ACTIVATION FUNCTIONS:
        activation_functions = nn.ModuleDict([
                                            ["Tanh", nn.Tanh()],
                                            ["ReLU", nn.ReLU()],
                                            ["LeakyReLU", nn.LeakyReLU()],
                                            ["Sigmoid", nn.Sigmoid()],
                                            ["Softmax", nn.Softmax()],
                                            ])
        # INPUT Linear function: 
        self.fc_in = nn.Linear(input_dim, hidden_dim)
        
        # INPUT Non-linearity (activation function): 
        self.tangh = nn.Tanh()
        
        # OUTPUT Linear function: 
        self.fc_out = nn.Linear(hidden_dim, output_dim)

        
    def forward(self, input):
        # INPUT Linear function:
        hidden = self.fc_in(input)

        # HIDDEN Non-linearity
        hidden = self.tangh(hidden)

        # OUTPUT Linear function
        output = self.fc_out(hidden)
        
        return output

model = FCN_2(2, 8, 1, activation="Tanh")
with torch.no_grad():
    weights_and_bias_initialization(model)
</code></pre>
<p>I get an ERROR related with <strong>nn.init.xavier_uniform_(module_with_layers.bias)</strong>:</p>
<pre><code class="lang-auto">ValueError: Fan in and fan out can not be computed for tensor with fewer than 2 dimensions.
</code></pre>
<p>Is the mistake related with the definition of the output?<br>
Thanks a lot</p>

soulitzer: <blockquote>
<p>I am not quite sure if a nn.ModuleDict is the best option for this case.</p>
</blockquote>
<p>That is correct, you probably just want to store the initialization if you prefer directly as an attribute on the module.</p>
<blockquote>
<p>I get an ERROR related with <strong>nn.init.xavier_uniform_(module_with_layers.bias)</strong></p>
</blockquote>
<p>xavier_uniform_ only works with 2D tensors, so for bias you would need to initialize it another way. One method that is commonly done is to initialize it with zeros_.</p>

luiscssz: <p>How can I store diferent type of initializations (Xavier_normal, Xavier_uniform_, Kaiming,) as attribute on the module and call it? Could you please add a code snippet?</p>
<p>Is there no possibility to use a typical initialization, Xavier_normal, Xavier_uniform_, Kaiming,… for the biases 1D?. I have read that sometimes the initialization with zeros is not the best option, because the model could bring you to the canonical solution, 0, as you starts from 0. Which another alternatives are available?</p>
<p>What is your opinion about the place of the initialization, see 2nd question above?</p>
<p>Thanks a lot</p>

soulitzer: <p>You can treat your module like any other Python subclass and store the callable as an attribute, e.g. do something like</p>
<pre data-code-wrap="python"><code class="lang-python">class FCN(nn.Module):
    def __init__(self, initialization: Callable):
          self.initialization = initialization
    def do_init(self):
          for module_with_layers in self.modules():
            if isinstance(module_with_layers, nn.Linear):
                self.initialization(module_with_layers.weight)

FCN(initialization=functools.partial(nn.init.xavier_uniform_, gain=nn.init.calculate_gain(...)))
</code></pre>
<blockquote>
<p>I have read that sometimes the initialization with zeros is not the best option, because the model could bring you to the canonical solution, 0, as you starts from 0.</p>
</blockquote>
<p>You shouldn’t be stuck at the canonical solution zero unless you also initialize weights to zero.</p>