Dropout.py

import torch.nn as nn
import sparseconvnet as scn
import torch.nn.functional as F

class Dropout(nn.Module):
    def __init__(self, p=0.5, inplace=False):
        super(Dropout, self).__init__()
        if p < 0 or p > 1:
            raise ValueError("dropout probability has to be between 0 and 1, "
                             "but got {}".format(p))
        self.p = p
        self.inplace = inplace

    def forward(self, input):
        output = SparseConvNetTensor()
        #F.dropout2d and F.dropout3d appear to be the same, but do not allow the 
        #symbolic method to make it an idenity op during inference mode (unlike F.dropout)
        output.features = F.dropout3d(input.features, self.p, self.training, self.inplace)
        output.metadata = input.metadata
        output.spatial_size = input.spatial_size
        return output

class Maxout(scn.MaxPooling):
    def __init__(self, dimension, pool_size, pool_stride, nFeatures, p=0.5):
        if p < 0 or p > 1:
            raise ValueError("dropout probability has to be between 0 and 1, "
                             "but got {}".format(p))

        super(Maxout, self).__init__(dimension, pool_size, pool_stride, nFeaturesToDrop=p*nFeatures)