braindecode.modules.TimmAttention#

class braindecode.modules.TimmAttention(dim, num_heads=8, qkv_bias=False, qv_bias=False, attn_drop=0.0, proj_drop=0.0, exact=False)[source]#

Multi-head self-attention with timm’s parameter names (qkv, proj), as in ViT/MAE/BEiT checkpoints. EEGPT and EEGDINO use it, so their released weights load with the same keys.

Parameters:

  • dim (int) – Token width.

  • num_heads (int) – Number of heads; head_dim = dim // num_heads.

  • qkv_bias (bool) – Bias of the fused qkv projection.

  • qv_bias (bool) – BEiT’s separate q_bias and v_bias parameters (the key has none) added after a bias-free qkv.

  • attn_drop (float) – Dropout on the attention weights.

  • proj_drop (float) – Dropout after the output projection.

  • exact (bool) – timm’s op order, softmax((q * scale) @ k^T) @ v, bit-identical to the reference implementations (see braindecode.functional.scaled_dot_product_attention()); otherwise torch’s kernels (Gaudi’s fused kernel on HPU).

Examples

>>> import torch
>>> from braindecode.modules import TimmAttention
>>> attn = TimmAttention(dim=16, num_heads=4)
>>> attn(torch.randn(2, 5, 16)).shape
torch.Size([2, 5, 16])

Methods

forward(x, need_weights=False)[source]#

(batch, seq, dim) -> (batch, seq, dim); with need_weights the attention weights (batch, heads, seq, seq) instead.

Return type:

Tensor