DOCSEARCHbyAlgolia
Submit Docs

deformable-attention

repository·main·Indexed 18 days ago

https://github.com/lucidrains/deformable-attention

A PyTorch implementation of Deformable Attention based on the 'Vision Transformer with Deformable Attention' paper. It provides DeformableAttention, DeformableAttention3D, and DeformableAttention1D classes to support 2D feature maps, volumetric/video inputs, and sequence data, incorporating Continuous Positional Embeddings from SwinV2 for improved extrapolation.

Tokens
1.2K
Snippets
4
Records
4
Agent score
13%

What's inside deformable-attention

  2. Use DeformableAttention for 2D inputs

    main

    The DeformableAttention class implements 2D deformable attention. It is designed for feature maps with dimensions (batch, dim, height, width).

    Key arguments:

    • dim: Feature dimensions.
    • dim_head: Dimension per head.
    • heads: Number of attention heads.
    • dropout: Dropout probability.
    • downsample_factor: Downsample factor (r in the paper).
    • offset_scale: Scale of offset, maximum offset.
    • offset_groups: Number of offset groups (should be a multiple of heads). If None, it defaults to standard behavior.
    • offset_kernel_size: Offset kernel size.
    import torch
from deformable_attention import DeformableAttention

attn = DeformableAttention(
    dim = 512,                   # feature dimensions
    dim_head = 64,               # dimension per head
    heads = 8,                   # attention heads
    dropout = 0.,                # dropout
    downsample_factor = 4,       # downsample factor (r in paper)
    offset_scale = 4,            # scale of offset, maximum offset
    offset_groups = None,        # number of offset groups, should be multiple of heads
    offset_kernel_size = 6,      # offset kernel size
)

x = torch.randn(1, 512, 64, 64)
attn(x) # (1, 512, 64, 64)

  3. Use DeformableAttention1D for sequence inputs

    main

    The DeformableAttention1D class implements 1D deformable attention for sequence data with dimensions (batch, dim, length).

    Key arguments:

    • dim: Feature dimensions.
    • downsample_factor: Downsample factor.
    • offset_scale: Scale of offset.
    • offset_kernel_size: Offset kernel size.
    import torch
from deformable_attention import DeformableAttention1D

attn = DeformableAttention1D(
    dim = 128,
    downsample_factor = 4,
    offset_scale = 2,
    offset_kernel_size = 6
)

x = torch.randn(1, 128, 512)
attn(x) # (1, 128, 512)

  4. Use DeformableAttention3D for volumetric/video inputs

    main

    The DeformableAttention3D class implements 3D deformable attention for inputs with dimensions (batch, dim, frames, height, width).

    Note that for 3D, several parameters accept tuples to specify values for each of the three spatial/temporal dimensions:

    • downsample_factor: A tuple representing the downsample factor for each dimension.
    • offset_scale: A tuple representing the scale of offset for each dimension.
    • offset_kernel_size: A tuple representing the offset kernel size for each dimension.
    import torch
from deformable_attention import DeformableAttention3D

attn = DeformableAttention3D(
    dim = 512,                          # feature dimensions
    dim_head = 64,                      # dimension per head
    heads = 8,                          # attention heads
    dropout = 0.,                       # dropout
    downsample_factor = (2, 8, 8),    # downsample factor (r in paper)
    offset_scale = (2, 8, 8),          # scale of offset, maximum offset
    offset_kernel_size = (4, 10, 10),   # offset kernel size
)

x = torch.randn(1, 512, 10, 32, 32) # (batch, dimension, frames, height, width)
attn(x) # (1, 512, 10, 32, 32)