LayerDropout

class torchtune.modules.LayerDropout(prob: float = 0.0, dim: Optional[int] = 0, disable_on_eval: Optional[bool] = True, seed: Optional[int] = None)

A module that applies layer dropout to the input tensor of an underlying module. It drops a portion of an input tensor, applies the underlying module on the remaining parts of the tensor, and then concatenates with the dropped portion of the tensor. When applied during training, it can have a regularization effect, and can potentially speedup training.

Parameters:

• prob (float) – The probability of dropping an input. Defaults to 0.0.

• dim (Optional[int]) – The dimension of input tensor along which to drop layers. Defaults to 0 (i.e., batch size).

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

• seed (Optional[int]) – The seed for the random number generator. Defaults to None.

Examples

>>> import torch
>>> # Apply layer dropout to a lambda function
>>> layer_dropout = LayerDropout(prob=0.5)
>>> output = layer_dropout(lambda x: x**2, torch.randn(1))
>>> # Apply layer dropout to a torch.nn.Linear module
>>> linear = torch.nn.Linear(5, 3)
>>> layer_dropout = LayerDropout(prob=0.5)
>>> output = layer_dropout(linear, torch.randn(1, 5))

forward(function: Union[Callable, Module], input: Tensor, *args, **kwargs) → Tensor

Apply layer dropout to the input tensor.

Parameters:

• function (Union[Callable, torch.nn.Module]) – The function or module to apply to the input tensor.

• input (torch.Tensor) – The input tensor.

• *args – Additional positional arguments passed to the function.

• **kwargs – Additional keyword arguments passed to the function.

Returns:

The output tensor after applying layer dropout.

Return type:

torch.Tensor