concrete_dropout.py

# usr/bin/env python3

# author: Conor McDonald

# torch==0.4.1
# numpy==1.14.3

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter
from torch import distributions
import numpy as np

class ConcreteDropout(nn.Module):
    """ Implementation of Concrete Dropout layer for efficient, 
    well calibrated uncertainty estimates. Enables approximation of Bayesian 
    deep learning via dropout. This means we can capture predictive uncertainty.
    See: Gal Y., Hron, J., Kendall, A. "Concrete Dropout", 2017.
    """

    def __init__(self, layer, input_shape, weight_reg=1e-6, 
                dropout_reg=1e-5, eps = 1e-7, temp = 0.1, **kwargs):
        super(ConcreteDropout).__init__(layer, *kwargs)
        self.layer = layer
        self.input_dim = np.prod(input_shape)
        self.eps = eps
        self.temp = temp

        self.weight_reg = weight_reg
        self.dropout_reg = dropout_reg

        self.uniform = distributions.Uniform(-2, 0)
        self.logit_p = nn.Parameter(self.uniform.sample())
        self.p = F.sigmoid(self.logit_p)

    def _concrete_dropout(self, x):
        
        unif_noise = self.uni.sample(x.shape)
        drop_prob = (
            (self.p + self.eps).log() 
            - (1. - self.p + self.eps)
            + (unif_noise + self.eps)
            - (1 - unif_noise + self.eps)
        )

        drop_prob = F.sigmoid(drop_prob + self.eps)
        random_tensor = 1 - drop_prob

        retain_prob = 1. - self.p
        x *= random_tensor
        x /= retain_prob
        return x
    
    def regularization(self):
        weights_reg = self.weight_reg * (self.layer.weight**2).sum() / (1 - self.p)
        dropout_reg = self.p * (self.p).log()
        dropout_reg += (1. - self.p) * (1 - self.p).log()
        dropout_reg *= self.dropout_reg * self.input_dim
        regularizer = (weights_reg + dropout_reg)
        return regularizer

    def forward(self, x):
        return self.layer(self._concrete_dropout(x))