Vector channel dropout where the elements of each vector channel are dropped together.
| 154 | |
| 155 | |
| 156 | class _VDropout(nn.Module): |
| 157 | ''' |
| 158 | Vector channel dropout where the elements of each |
| 159 | vector channel are dropped together. |
| 160 | ''' |
| 161 | def __init__(self, drop_rate): |
| 162 | super().__init__() |
| 163 | self.drop_rate = drop_rate |
| 164 | self.dummy_param = nn.Parameter(torch.empty(0)) |
| 165 | |
| 166 | def forward(self, x): |
| 167 | ''' |
| 168 | :param x: `torch.Tensor` corresponding to vector channels |
| 169 | ''' |
| 170 | device = self.dummy_param.device |
| 171 | if not self.training: |
| 172 | return x |
| 173 | mask = torch.bernoulli( |
| 174 | (1 - self.drop_rate) * torch.ones(x.shape[:-1], device=device) |
| 175 | ).unsqueeze(-1) |
| 176 | x = mask * x / (1 - self.drop_rate) |
| 177 | return x |
| 178 | |
| 179 | |
| 180 | class Dropout(nn.Module): |