| | |
| | |
| | |
| | |
| | """ |
| | LayerDrop as described in https://arxiv.org/abs/1909.11556. |
| | """ |
| |
|
| | import torch |
| | import torch.nn as nn |
| |
|
| |
|
| | class LayerDropModuleList(nn.ModuleList): |
| | """ |
| | A LayerDrop implementation based on :class:`torch.nn.ModuleList`. |
| | |
| | We refresh the choice of which layers to drop every time we iterate |
| | over the LayerDropModuleList instance. During evaluation we always |
| | iterate over all layers. |
| | |
| | Usage:: |
| | |
| | layers = LayerDropList(p=0.5, modules=[layer1, layer2, layer3]) |
| | for layer in layers: # this might iterate over layers 1 and 3 |
| | x = layer(x) |
| | for layer in layers: # this might iterate over all layers |
| | x = layer(x) |
| | for layer in layers: # this might not iterate over any layers |
| | x = layer(x) |
| | |
| | Args: |
| | p (float): probability of dropping out each layer |
| | modules (iterable, optional): an iterable of modules to add |
| | """ |
| |
|
| | def __init__(self, p, modules=None): |
| | super().__init__(modules) |
| | self.p = p |
| |
|
| | def __iter__(self): |
| | dropout_probs = torch.empty(len(self)).uniform_() |
| | for i, m in enumerate(super().__iter__()): |
| | if not self.training or (dropout_probs[i] > self.p): |
| | yield m |
| |
|
