prepare_layer_dropout

torchtune.modules.prepare_layer_dropout(layers: Union[ModuleList, Iterable[Module]], prob_max: float = 0.0, prob_layer_scale: Optional[ScaleType] = ScaleType.UNIFORM, layers_str: Optional[str] = None, disable_on_eval: Optional[bool] = True) → None

Prepare a model’s layers for layer dropout by wrapping each layer with a ModuleLayerDropoutWrapper. This function takes in a list of layers, the maximum probability of dropping a layer, the scaling type for the layer dropout probability, a string specifying which layers to apply dropout to, and a boolean indicating whether to disable dropout during evaluation. It then wraps each layer of the model inplace with a ModuleLayerDropoutWrapper, which applies layer dropout to the input tensor.

Parameters:

• layers (Union[torch.nn.ModuleList, Iterable[torch.nn.Module]]) – The list of layers to prepare for layer dropout.

• prob_max (float) – The maximum probability of dropping a layer. Defaults to 0.0.

• prob_layer_scale (Optional[ScaleType]) – The scaling type for the dropout probability across layers. Defaults to ScaleType.UNIFORM.

• layers_str (Optional[str]) – A string specifying which layers to apply dropout to. Defaults to None which means apply to all layers.

• disable_on_eval (Optional[bool]) – Whether to disable dropout during evaluation. Defaults to True.

Returns:

None

Example

>>> import torch
>>> from torch import nn
>>> # Define a simple model
>>> class MyModel(nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.layers = nn.ModuleList([
...             nn.Linear(5, 3),
...             nn.Linear(3, 2),
...             nn.Linear(2, 1),
...             nn.Linear(1, 2),
...             nn.Linear(2, 3),
...         ])
...
...     def forward(self, x):
...         for layer in self.layers:
...             x = layer(x)
...         return x
>>> model = MyModel()
>>> # Apply layer dropout uniformly to all layers
>>> prepare_layer_dropout(model.layers, prob_max=0.2, prob_layer_scale=ScaleType.UNIFORM)
>>> # Apply layer dropout every other layer, as described in LayerDrop paper
    (Fan et al., https://arxiv.org/abs/1909.11556v1)
>>> prepare_layer_dropout(model.layers, prob_max=0.2, prob_layer_scale=ScaleType.UNIFORM, layers_str="::2")
>>> # Apply layer dropout that increases linearly across layers, as described in Progressive Layer
    Dropout paper (Zhang et al., https://arxiv.org/abs/2010.13369)
>>> prepare_layer_dropout(model.layers, prob_max=0.2, prob_layer_scale=ScaleType.LINEAR)
>>> # Apply layer dropout that increases exponentially across layers, as described in
    LayerSkip paper (Elhoushi et al., https://arxiv.org/abs/2404.16710)
>>> prepare_layer_dropout(model.layers, prob_max=0.2, prob_layer_scale=ScaleType.EXP)