mmlearn.modules.layers.patch_dropout

Patch dropout layer.

Classes

PatchDropout

Patch dropout layer.

class PatchDropout(keep_rate=0.5, bias=None, token_shuffling=False)

Patch dropout layer.

Drops patch tokens (after embedding and adding CLS token) from the input tensor. Usually used in vision transformers to reduce the number of tokens. [1]

Parameters:

• keep_rate (float, optional, default=0.5) – The proportion of tokens to keep.

• bias (Optional[float], optional, default=None) – The bias to add to the random noise before sorting.

• token_shuffling (bool, optional, default=False) – If True, the tokens are shuffled.

References

[1]

Liu, Y., Matsoukas, C., Strand, F., Azizpour, H., & Smith, K. (2023). Patchdropout: Economizing vision transformers using patch dropout. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (pp. 3953-3962).

forward(x, force_drop=False)

Drop tokens from the input tensor.

Parameters:

• x (torch.Tensor) – Input tensor of shape (batch_sz, seq_len, dim).

• force_drop (bool, optional, default=False) – If True, the tokens are always dropped, even when the model is in evaluation mode.

Returns:

Tensor of shape (batch_sz, keep_len, dim) containing the kept tokens.

Return type:

torch.Tensor

uniform_mask(x)

Generate token ids to keep from uniform random distribution.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_sz, seq_len, dim).

Returns:

Tensor of shape (batch_sz, keep_len) containing the token ids to keep.

Return type:

torch.Tensor