DOCSEARCHbyAlgolia
Submit Docs

RIN-pytorch: Recurrent Interface Network

repository·main·Indexed 42 days ago

https://github.com/lucidrains/recurrent-interface-network-pytorch

A PyTorch implementation of the Recurrent Interface Network (RIN) for efficient image and video generation. Version 0.8.0 features a specialized sigmoid-based noise schedule, latent self-conditioning, and recurrent steps to avoid cascading networks. Includes the RIN model class, GaussianDiffusion wrapper for the diffusion process, and a Trainer class for high-level API training.

Tokens
1.2K
Snippets
5
Records
5
Agent score
3%

What's inside RIN-pytorch

  2. Manual training and sampling loop

    main

    You can use RIN and GaussianDiffusion directly without the Trainer class for custom training loops or inference.

    Note: Training images should be normalized from 0 to 1. After training, use .sample() to generate new images.

    import torch
from rin_pytorch import RIN, GaussianDiffusion

model = RIN(
    dim = 256,
    image_size = 128,
    patch_size = 8,
    depth = 6,
    num_latents = 128,
    dim_latent = 512,
    latent_self_attn_depth = 4,
).cuda()

diffusion = GaussianDiffusion(
    model,
    timesteps = 1000,
    train_prob_self_cond = 0.9,
    scale = 1.
).cuda()

# Training step
training_images = torch.randn(8, 3, 128, 128).cuda() # images are normalized from 0 to 1
loss = diffusion(training_images)
loss.backward()

# Sampling step (after training)
sampled_images = diffusion.sample(batch_size = 4)
sampled_images.shape # (4, 3, 128, 128)

  3. Train using the Trainer class

    main

    The Trainer class provides a high-level API for training the model on a dataset of images. Results are saved periodically to a ./results folder.

    Key arguments:

    • diffusion: The GaussianDiffusion instance.
    • path: Path to your image dataset.
    • num_samples: Number of samples to generate for evaluation.
    • train_batch_size: Batch size for training.
    • gradient_accumulate_every: Number of steps to accumulate gradients.
    • train_lr: Learning rate.
    • save_and_sample_every: Frequency of saving checkpoints and samples.
    • train_num_steps: Total number of training steps.
    • ema_decay: Exponential moving average decay.
    from rin_pytorch import RIN, GaussianDiffusion, Trainer

# ... setup model and diffusion ...

trainer = Trainer(
    diffusion,
    '/path/to/your/images',
    num_samples = 16,
    train_batch_size = 4,
    gradient_accumulate_every = 4,
    train_lr = 1e-4,
    save_and_sample_every = 1000,
    train_num_steps = 700000,
    ema_decay = 0.995,
)

trainer.train()

  4. Configure GaussianDiffusion

    main

    The GaussianDiffusion class wraps the RIN model to handle the diffusion process.

    Key arguments:

    • model: An instance of the RIN class.
    • timesteps: Number of diffusion timesteps.
    • train_prob_self_cond: Probability of self-conditioning on latents during training.
    • scale: Used to noise higher resolution images more. Setting this to a value < 1. leads to better convergence when training on high-resolution images (e.g., 512, 1024). Input noised images are automatically variance normalized.
    from rin_pytorch import GaussianDiffusion

diffusion = GaussianDiffusion(
    model,
    timesteps = 400,
    train_prob_self_cond = 0.9,
    scale = 1.
)

  5. Initialize the RIN model

    main

    The RIN class implements the Recurrent Interface Network. It requires parameters for dimensions, image size, and latent configuration.

    Key arguments:

    • dim: Model dimensions.
    • image_size: Target image size.
    • patch_size: Size of the patches.
    • depth: Network depth.
    • num_latents: Number of latents (the paper uses 256).
    • dim_latent: Dimension of the latents (can be larger than dim for higher capacity).
    • latent_self_attn_depth: Number of latent self-attention blocks per recurrent step (K in the paper).
    from rin_pytorch import RIN

model = RIN(
    dim = 256,
    image_size = 128,
    patch_size = 8,
    depth = 6,
    num_latents = 128,
    dim_latent = 512,
    latent_self_attn_depth = 4,
)