Bayesian Dropout Layer

BayesianDropout.py

# Bayesian dropout layer: Applies random or optionally pre defined dropout to input during both training and evaluation
class BayesianDropout(torch.nn.Module):

    def __init__(self, drop_rate):
        super(BayesianDropout, self).__init__()
        self.p = drop_rate

    def forward(self, x, mask=None):
        if mask is None:
            return F.dropout(x, self.p, training=True, inplace=False)
        else:
            return x * mask
         
# Example usage:
b_dropout = BayesianDropout(drop_rate=0.33)