DOCSEARCHbyAlgolia
Submit Docs

denoising-diffusion-pytorch

repository·main·Indexed 27 days ago

https://github.com/lucidrains/denoising-diffusion-pytorch

A PyTorch implementation of Denoising Diffusion Probabilistic Models (DDPM) for generative modeling. It supports both 2D images and 1D sequences, providing components such as Unet, GaussianDiffusion, and a Trainer class that handles batching, EMA, and mixed precision. It also includes support for multi-GPU training via 🤗 Accelerator.

Tokens
1.5K
Snippets
5
Records
5
Agent score
45%

What's inside denoising-diffusion-pytorch

  2. Configure Multi-GPU Training with Accelerate

    main

    The Trainer class supports 🤗 Accelerator. To perform multi-GPU training, use the accelerate CLI in your project root:

    1. Run accelerate config to configure your environment.
    2. Run accelerate launch <your_training_script>.py to start training.
    $ accelerate config

$ accelerate launch train.py

  3. Basic Usage with Unet and GaussianDiffusion

    main

    To implement a standard 2D diffusion model, initialize a Unet and wrap it in a GaussianDiffusion object. During training, pass normalized images (0 to 1) to the diffusion object to compute the loss. After training, use the .sample() method to generate new images.

    import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion

model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8),
    flash_attn = True
)

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000    # number of steps
)

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

# after a lot of training

sampled_images = diffusion.sample(batch_size = 4)
sampled_images.shape # (4, 3, 128, 128)

  4. Train a model using the Trainer class

    main

    The Trainer class simplifies training by allowing you to point to an image folder. It handles batching, learning rate, gradient accumulation, EMA, and mixed precision (AMP). Samples and checkpoints are automatically logged to ./results.

    from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer

model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8),
    flash_attn = True
)

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,           # number of steps
    sampling_timesteps = 250    # number of sampling timesteps (using ddim for faster inference [see citation for ddim paper])
)

trainer = Trainer(
    diffusion,
    'path/to/your/images',
    train_batch_size = 32,
    train_lr = 8e-5,
    train_num_steps = 700000,         # total training steps
    gradient_accumulate_every = 2,    # gradient accumulation steps
    ema_decay = 0.995,                # exponential moving average decay
    amp = True,                       # turn on mixed precision
    calculate_fid = True              # whether to calculate fid during training
)

trainer.train()

  5. Use 1D Sequence Diffusion (Unet1D)

    main

    For 1D sequence data, use the Unet1D, GaussianDiffusion1D, Trainer1D, and Dataset1D classes. Ensure your input features are normalized from 0 to 1. Note that Trainer1D does not perform automatic sample evaluation as the data type is unknown.

    import torch
from denoising_diffusion_pytorch import Unet1D, GaussianDiffusion1D, Trainer1D, Dataset1D

model = Unet1D(
    dim = 64,
    dim_mults = (1, 2, 4, 8),
    channels = 32
)

diffusion = GaussianDiffusion1D(
    model,
    seq_length = 128,
    timesteps = 1000,
    objective = 'pred_v'
)

training_seq = torch.rand(64, 32, 128) # features are normalized from 0 to 1

loss = diffusion(training_seq)
loss.backward()

# Or using trainer

dataset = Dataset1D(training_seq)  # this is just an example, but you can formulate your own Dataset and pass it into the `Trainer1D` below

trainer = Trainer1D(
    diffusion,
    dataset = dataset,
    train_batch_size = 32,
    train_lr = 8e-5,
    train_num_steps = 700000,
    gradient_accumulate_every = 2,
    ema_decay = 0.995,
    amp = True,
)
trainer.train()

# after a lot of training

sampled_seq = diffusion.sample(batch_size = 4)
sampled_seq.shape # (4, 32, 128)