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, modu
| 137 | |
| 138 | # modify from : https://github.com/facebookresearch/fairseq/blob/main/fairseq/modules/layer_drop.py # noqa |
| 139 | class LayerDropModuleList(torch.nn.ModuleList): |
| 140 | """ |
| 141 | A LayerDrop implementation based on :class:`torch.nn.ModuleList`. |
| 142 | |
| 143 | We refresh the choice of which layers to drop every time we iterate |
| 144 | over the LayerDropModuleList instance. During evaluation we always |
| 145 | iterate over all layers. |
| 146 | |
| 147 | Usage:: |
| 148 | |
| 149 | layers = LayerDropList(p=0.5, modules=[layer1, layer2, layer3]) |
| 150 | for layer in layers: # this might iterate over layers 1 and 3 |
| 151 | x = layer(x) |
| 152 | for layer in layers: # this might iterate over all layers |
| 153 | x = layer(x) |
| 154 | for layer in layers: # this might not iterate over any layers |
| 155 | x = layer(x) |
| 156 | |
| 157 | Args: |
| 158 | p (float): probability of dropping out each layer |
| 159 | modules (iterable, optional): an iterable of modules to add |
| 160 | |
| 161 | Limitations: |
| 162 | 1 can work with ddp when layer's gradient checkpoint disabled |
| 163 | 2 can't work with ddp when layer's gradient checkpoint enables |
| 164 | 3 can work with fsdp |
| 165 | 4 can work with deepspeed |
| 166 | """ |
| 167 | |
| 168 | def __init__(self, p: List[float], modules=None): |
| 169 | super().__init__(modules) |
| 170 | assert len(p) == len(self) |
| 171 | self.p = p |
| 172 | |
| 173 | def __iter__(self): |
| 174 | dropout_probs = torch.empty(len(self)).uniform_() |
| 175 | for i, m in enumerate(super().__iter__()): |
| 176 | if not self.training or (dropout_probs[i] > self.p[i]): |
| 177 | yield m |