https://github.com/lucidrains/bit-diffusion

Implementation of Bit Diffusion, Hinton's group's attempt at discrete denoising diffusion, in Pytorch
https://github.com/lucidrains/bit-diffusion

artificial-intelligence deep-learning denoising-diffusion discrete

Last synced: 7 days ago
JSON representation

Implementation of Bit Diffusion, Hinton's group's attempt at discrete denoising diffusion, in Pytorch

• Host: GitHub
• URL: https://github.com/lucidrains/bit-diffusion
• Owner: lucidrains
• License: mit
• Created: 2022-08-17T15:41:22.000Z (over 2 years ago)
• Default Branch: main
• Last Pushed: 2023-10-14T14:43:45.000Z (about 1 year ago)
• Last Synced: 2024-10-30T06:58:27.751Z (21 days ago)
• Topics: artificial-intelligence, deep-learning, denoising-diffusion, discrete
• Language: Python
• Homepage:
• Size: 75.2 KB
• Stars: 332
• Watchers: 5
• Forks: 17
• Open Issues: 4
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

        


## Bit Diffusion - Pytorch

Implementation of Bit Diffusion, Hinton's group's attempt at discrete denoising diffusion, in Pytorch

It seems like they missed the mark for text, but the research direction still seems promising. I think a clean repository will do the research community a lot of benefits for those branching off from here.

## Install

```bash

$ pip install bit-diffusion

```

## Usage

```python

from bit_diffusion import Unet, Trainer, BitDiffusion

model = Unet(

    dim = 32,

    channels = 3,

    dim_mults = (1, 2, 4, 8),

).cuda()

bit_diffusion = BitDiffusion(

    model,

    image_size = 128,

    timesteps = 100,

    time_difference = 0.1,       # they found in the paper that at lower number of timesteps, a time difference during sampling of greater than 0 helps FID. as timesteps increases, this time difference can be set to 0 as it does not help

    use_ddim = True              # use ddim

).cuda()

trainer = Trainer(

    bit_diffusion,

    '/path/to/your/data',             # path to your folder of images

    results_folder = './results',     # where to save results

    num_samples = 16,                 # number of samples

    train_batch_size = 4,             # training batch size

    gradient_accumulate_every = 4,    # gradient accumulation

    train_lr = 1e-4,                  # learning rate

    save_and_sample_every = 1000,     # how often to save and sample

    train_num_steps = 700000,         # total training steps

    ema_decay = 0.995,                # exponential moving average decay

)

trainer.train()

```

Results will be saved periodically to the `./results` folder

If you would like to experiment with the `Unet` and `BitDiffusion` class outside the `Trainer`

```python

import torch

from bit_diffusion import Unet, BitDiffusion

model = Unet(

    dim = 64,

    dim_mults = (1, 2, 4, 8)

)

bit_diffusion = BitDiffusion(

    model,

    image_size = 128,

    timesteps = 1000

)

training_images = torch.randn(8, 3, 128, 128) # images are normalized from 0 to 1

loss = bit_diffusion(training_images)

loss.backward()

# after a lot of training

sampled_images = bit_diffusion.sample(batch_size = 4)

sampled_images.shape # (4, 3, 128, 128)

```

## Citations

```bibtex

@article{Chen2022AnalogBG,

    title   = {Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning},

    author  = {Ting Chen and Ruixiang Zhang and Geoffrey E. Hinton},

    journal = {ArXiv},

    year    = {2022},

    volume  = {abs/2208.04202}

}

```