DOCSEARCHbyAlgolia
Submit Docs

magvit2-pytorch

repository·main·Indexed 20 days ago

https://github.com/lucidrains/magvit2-pytorch

A PyTorch implementation of the MagViT2 architecture for video generation and understanding. It provides a VideoTokenizer for advanced video tokenization and a VideoTokenizerTrainer for training on video or image datasets. Features include support for Exponential Moving Average (EMA) tokenization for inference and experiment tracking via Weights & Biases.

Tokens
1.2K
Snippets
4
Records
4
Agent score
21%

What's inside magvit2-pytorch

  2. Initialize and train a VideoTokenizer

    main

    To use MagViT2, you can instantiate a VideoTokenizer with a specific architecture configuration and then use the VideoTokenizerTrainer to train it on a dataset of videos or images.

    Key parameters for VideoTokenizer:

    • image_size: The resolution of the input.
    • init_dim: Initial dimension.
    • max_dim: Maximum dimension.
    • codebook_size: Size of the codebook.
    • layers: A tuple defining the architecture (e.g., 'residual', 'compress_space', 'compress_time', 'attend_space', 'attend_time').

    Key parameters for VideoTokenizerTrainer:

    • tokenizer: The VideoTokenizer instance.
    • dataset_folder: Path to the media folder.
    • dataset_type: Either 'videos' or 'images'.
    • batch_size: Training batch size.
    • grad_accum_every: Gradient accumulation steps.
    • learning_rate: Optimizer learning rate.
    • num_train_steps: Total training steps.
    from magvit2_pytorch import (
    VideoTokenizer,
    VideoTokenizerTrainer
)

tokenizer = VideoTokenizer(
    image_size = 128,
    init_dim = 64,
    max_dim = 512,
    codebook_size = 1024,
    layers = (
        'residual',
        'compress_space',
        ('consecutive_residual', 2),
        'compress_space',
        ('consecutive_residual', 2),
        'linear_attend_space',
        'compress_space',
        ('consecutive_residual', 2),
        'attend_space',
        'compress_time',
        ('consecutive_residual', 2),
        'compress_time',
        ('consecutive_residual', 2),
        'attend_time',
    )
)

trainer = VideoTokenizerTrainer(
    tokenizer,
    dataset_folder = '/path/to/a/lot/of/media',
    dataset_type = 'videos',
    batch_size = 4,
    grad_accum_every = 8,
    learning_rate = 2e-5,
    num_train_steps = 1_000_000
)

trainer.train()

  3. Tokenize and decode video using the EMA tokenizer

    main

    After training, you can use the Exponential Moving Average (EMA) version of the tokenizer for inference. This allows you to convert video tensors into discrete code indices and reconstruct the video from those indices.

    • ema_tokenizer.tokenize(video): Converts a video tensor of shape (B, C, T, H, W) into discrete codes.
    • ema_tokenizer.decode_from_code_indices(codes): Reconstructs the video from the discrete codes.
    • ema_tokenizer(video, return_recon=True): Performs a standard forward pass to return the reconstructed video.
    # Assuming trainer has been trained
ema_tokenizer = trainer.ema_tokenizer

# Mock video: (Batch, Channels, Time, Height, Width)
video = torch.randn(1, 3, 17, 128, 128)

# Tokenizing video to discrete codes
codes = ema_tokenizer.tokenize(video) 

# Decoding back to video
decoded_video = ema_tokenizer.decode_from_code_indices(codes)

# Sanity check
assert torch.allclose(
    decoded_video,
    ema_tokenizer(video, return_recon = True)
)

  4. Track experiments with Weights & Biases

    main

    To enable experiment tracking via Weights & Biases (W&B), set use_wandb_tracking = True in the VideoTokenizerTrainer constructor. Use the .trackers context manager to define the project and run names during training.

    trainer = VideoTokenizerTrainer(
    use_wandb_tracking = True,
    ...
)

with trainer.trackers(project_name = 'magvit2', run_name = 'baseline'):
    trainer.train()