DOCSEARCHbyAlgolia
Submit Docs

mlp-mixer-pytorch

repository·main·Indexed 21 days ago

https://github.com/lucidrains/mlp-mixer-pytorch

A PyTorch implementation of the MLP-Mixer architecture for computer vision, providing an all-MLP solution that avoids convolutions and attention mechanisms. Includes the MLPMixer class for 2D images and the MLPMixer3D class for video data processing.

Tokens
697
Snippets
3
Records
3
Agent score
27%

What's inside mlp-mixer-pytorch

  2. Use MLPMixer3D for video data

    main

    The MLPMixer3D class extends the MLP-Mixer architecture to handle temporal dimensions for video processing.

    Parameters:

    • image_size: Tuple representing spatial dimensions (height, width).
    • time_size: Number of frames in the temporal dimension.
    • time_patch_size: Size of the temporal patches.
    • channels: Number of input channels.
    • patch_size: Spatial patch size.
    • dim: Embedding dimension.
    • depth: Number of mixer blocks.
    • num_classes: Number of output classes.
    import torch
from mlp_mixer_pytorch import MLPMixer3D

model = MLPMixer3D(
    image_size = (256, 128),
    time_size = 4,
    time_patch_size = 2,
    channels = 3,
    patch_size = 16,
    dim = 512,
    depth = 12,
    num_classes = 1000
)

video = torch.randn(1, 3, 4, 256, 128)
pred = model(video) # (1, 1000)

  3. Use MLPMixer for 2D images

    main

    The MLPMixer class implements an all-MLP architecture for vision tasks. It supports both square and rectangular images.

    Parameters:

    • image_size: Integer (for square) or tuple (for rectangular) representing image dimensions.
    • channels: Number of input channels (e.g., 3 for RGB).
    • patch_size: Size of the patches used for tokenization.
    • dim: Embedding dimension.
    • depth: Number of mixer blocks.
    • num_classes: Number of output classes.
    import torch
from mlp_mixer_pytorch import MLPMixer

# Square image example
model = MLPMixer(
    image_size = 256,
    channels = 3,
    patch_size = 16,
    dim = 512,
    depth = 12,
    num_classes = 1000
)

img = torch.randn(1, 3, 256, 256)
pred = model(img) # (1, 1000)

# Rectangular image example
model = MLPMixer(
    image_size = (256, 128),
    channels = 3,
    patch_size = 16,
    dim = 512,
    depth = 12,
    num_classes = 1000
)

img = torch.randn(1, 3, 256, 128)
pred = model(img) # (1, 1000)