How to add dropout layers automatically to a neural network in pytorch

Question

I have a neural network in pytorch and make each layer automatically via the following structure:

class FCN(nn.Module):
    ##Neural Network
    def __init__(self,layers):
        super().__init__() #call __init__ from parent class 
        self.activation = nn.Tanh()
        self.loss_function = nn.MSELoss(reduction ='mean')
        'Initialise neural network as a list using nn.Modulelist'  
        self.linears = nn.ModuleList([nn.Linear(layers[i], layers[i+1]) for i in range(len(layers)-1)])
        self.iter = 0
        'Xavier Normal Initialization'
        for i in range(len(layers)-1):         
            nn.init.xavier_normal_(self.linears[i].weight.data, gain=1.0)            
            nn.init.zeros_(self.linears[i].bias.data)   
    'foward pass'
    def forward(self, x):
        if torch.is_tensor(x) != True:         
            x = torch.from_numpy(x)                
        a = x.float()
        for i in range(len(layers)-2):  
            z = self.linears[i](a)              
            a = self.activation(z)    
        a = self.linears[-1](a)
        return a

The following code also makes the network for me:

layers = np.array([2, 50, 50, 1])
model = FCN(layers)

Now, I am wondering how I can automatically add dropout layers to the network. I tried the following change in the network structure but it only gives me one dropout layer at the end:

self.linears = nn.ModuleList([nn.Linear(layers[i], layers[i+1]) for i in range(len(layers)-1) + nn.Dropout(p=0.5)]

I very much appreciate any help in this regard.

Answer 1

If you can add a dropout layer by "adding it" with + as you do (I havent seen that, but if it works that is dope!) you should just move the + DropOut before the range I assume i.e

self.linears = nn.ModuleList([nn.Linear(layers[i], layers[i+1])+ nn.Dropout(p=0.5) for i in range(len(layers)-1) ]

EDIT

As expected you can't add it like that.

What you would do is to add a list with dropout-layers in the same way you do linear-layers, which you then use in your forward pass.

Below is an example; it might need to be tweaked to match your inputs etc

class FCN(nn.Module):
    ## Neural Network
    def __init__(self,layers):
        super().__init__()
        self.activation = nn.Tanh()
        self.loss_function = nn.MSELoss(reduction ='mean')
        'Initialise neural network as a list using nn.Modulelist'  
        self.linears = nn.ModuleList([nn.Linear(layers[i], layers[i+1]) for i in range(len(layers)-1)]) 
        self.dropout_layers = [nn.Dropout(p=0.5) for _ in range(len(layers)-1)]
        self.iter = 0
        'Xavier Normal Initialization'
        for i in range(len(layers)-1):         
            nn.init.xavier_normal_(self.linears[i].weight.data, gain=1.0)            
            nn.init.zeros_(self.linears[i].bias.data)

        def forward(self,x):
            for layer,dropout in zip(self.linears,self.dropout_layers):
                 x = layer(x)
                 x = dropout(x)
            return x

Answer 2

I think you should rewrite your code, nn.Sequential() might be the best tool to use.

class FCN(nn.Module):
    """Neural Network"""
    def __init__(self, layers, drop_p=0.5):
        super().__init__()  # call __init__ from parent class
        self.loss_function = nn.MSELoss(reduction='mean')

        # Use nn.Sequential to create the neural network
        # Here you need a list of Module that you want to use
        module_list = [[nn.Linear(layers[i], layers[i + 1]), nn.Dropout(p=drop_p), nn.Tanh()] for i in range(len(layers) - 2)]
        self.linears = nn.Sequential(*[item for sublist in module_list for item in sublist])
        self.last_layer = nn.Linear(layers[-2], layers[-1])

        self.iter = 0
        # Xavier Normal Initialization
        self.linears.apply(self.weights_init)
        self.last_layer.apply(self.weights_init)

    @staticmethod
    def weights_init(m):
        # Xavier Normal Initialization
        if isinstance(m, nn.Conv2d):
            torch.nn.init.xavier_uniform_(m.weight)
            torch.nn.init.zero_(m.bias)

    # Forward pass
    def forward(self, x):
        if not torch.is_tensor(x):
            x = torch.from_numpy(x)

        x = x.float()
        x = self.linears(x)
        x = self.last_layer(x)
        return x

Something like this works fine, you can test it out with a dummy tensor like:

if __name__ == "__main__":
    layers_size = np.array([2, 50, 50, 1])
    model = FCN(layers_size)

    t = torch.rand((2,))
    output = model(t)
    print(output)

>>> tensor([-0.0045], grad_fn=<AddBackward0>)

For the flattening operation of the module_list variable you can check here.