dropout.py

import torch
import torch.nn
generator = torch.Generator(device="cuda")
generator.manual_seed(42)
def createMask(x, dr):
    mask = x.new().resize_as_(x).bernoulli_(1 - dr,generator=generator).div_(1 - dr).detach_()
    # print('droprate='+str(dr))
    return mask


class DropMask(torch.autograd.function.InplaceFunction):
    @classmethod
    def forward(cls, ctx, input, mask, train=False, inplace=False):
        ctx.train = train
        ctx.inplace = inplace
        ctx.mask = mask

        if not ctx.train:
            return input
        else:
            if ctx.inplace:
                ctx.mark_dirty(input)
                output = input
            else:
                output = input.clone()
            output.mul_(ctx.mask)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        if ctx.train:
            return grad_output * ctx.mask, None, None, None
        else:
            return grad_output, None, None, None