<filename>torchnn/models/variational_dropout.py<gh_stars>0
import torch
import torch.nn as nn
from torch.nn.utils.rnn import PackedSequence

class VariationalDropout(nn.Module):
    def __init__(self, dropout, batch_first):
        super().__init__()
        self.dropout = dropout
        self.batch_first = batch_first

    def forward(self, x):
        if not self.training or self.dropout <= 0.:
            return x

        is_packed = isinstance(x, PackedSequence)
        if is_packed:
            x, batch_sizes = x
            max_batch_size = int(batch_sizes[0])
        else:
            batch_sizes = None
            max_batch_size = x.size(0)

        # Drop same mask across entire sequence
        if self.batch_first:
            m = x.new_empty(max_batch_size, 1, x.size(2), requires_grad=False).bernoulli_(1 - self.dropout)
        else:
            m = x.new_empty(1, max_batch_size, x.size(2), requires_grad=False).bernoulli_(1 - self.dropout)
        x = x.masked_fill(m == 0, 0) / (1 - self.dropout)

        if is_packed:
            return PackedSequence(x, batch_sizes)
        else:
            return x

<filename>torchnn/models/variational_dropout.py<gh_stars>0
import torch
import torch.nn as nn
from torch.nn.utils.rnn import PackedSequence

class VariationalDropout(nn.Module):
    def __init__(self, dropout, batch_first):
        super().__init__()
        self.dropout = dropout
        self.batch_first = batch_first

    def forward(self, x):
        if not self.training or self.dropout <= 0.:
            return x

        is_packed = isinstance(x, PackedSequence)
        if is_packed:
            x, batch_sizes = x
            max_batch_size = int(batch_sizes[0])
        else:
            batch_sizes = None
            max_batch_size = x.size(0)

        # Drop same mask across entire sequence
        if self.batch_first:
            m = x.new_empty(max_batch_size, 1, x.size(2), requires_grad=False).bernoulli_(1 - self.dropout)
        else:
            m = x.new_empty(1, max_batch_size, x.size(2), requires_grad=False).bernoulli_(1 - self.dropout)
        x = x.masked_fill(m == 0, 0) / (1 - self.dropout)

        if is_packed:
            return PackedSequence(x, batch_sizes)
        else:
            return x

# Drop same mask across entire sequence