| 63 | return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) |
| 64 | |
| 65 | class Dropout(nn.Module): |
| 66 | # Dropout entire row or column |
| 67 | def __init__(self, broadcast_dim=None, p_drop=0.15): |
| 68 | super(Dropout, self).__init__() |
| 69 | # give ones with probability of 1-p_drop / zeros with p_drop |
| 70 | self.sampler = torch.distributions.bernoulli.Bernoulli(torch.tensor([1-p_drop])) |
| 71 | self.broadcast_dim=broadcast_dim |
| 72 | self.p_drop=p_drop |
| 73 | def forward(self, x): |
| 74 | if not self.training: # no drophead during evaluation mode |
| 75 | return x |
| 76 | shape = list(x.shape) |
| 77 | if not self.broadcast_dim == None: |
| 78 | shape[self.broadcast_dim] = 1 |
| 79 | mask = self.sampler.sample(shape).to(x.device).view(shape) |
| 80 | |
| 81 | x = mask * x / (1.0 - self.p_drop) |
| 82 | return x |
| 83 | |
| 84 | def rbf(D): |
| 85 | # Distance radial basis function |