Custom dropout in PyTorch?

Question

I’m looking to implement a custom dropout in PyTorch — in the sense that I’d like to pass a mask of some sort, and have the corresponding neurons be “dropped out”, rather than dropping out random neurons with a specific probability. Something like:

Nn.Dropout(inputs, mask = [0, 1, 0, …]

I can’t seem to find anything in pytorch’s documentation or forums, so any and all help would be appreciated.

Answer 1

Standard dropout samples a random mask with every forward pass. This is a basic example:

class Dropout(nn.Module):
    def __init__(self, p):
        super().__init__()
        self.p = p
        
    def get_mask(self, x):
        mask = torch.rand(*x.shape)<=self.p
        return mask
        
    def forward(self, x):
        if self.training:
            mask = self.get_mask(x)
            x = x * mask
        return x

If you want a custom dropout, you can implement your own logic in get_mask