MegEngine implementation of Diffusion Models
Project description
MegDiffusion
MegEngine implementation of Diffusion Models (in early development).
Current maintainer: @MegChai
Usage
Infer with pre-trained models
Now users can use megengine.hub
to get pre-trained models directly:
megengine.hub.list("MegEngine/MegDiffusion:main")
megengine.hub.help("MegEngine/MegDiffusion:main", "ddpm_cifar10")
model = megengine.hub.load("MegEngine/MegDiffusion:main", "ddpm_cifar10", pretrained=True)
model.eval()
Or if you have downloaded or installed MegDiffusion, you can get pre-trained models from model
module.
from megdiffusion.model import ddpm_cifar10
model = ddpm_cifar10(pretrained=True)
model.eval()
The inference script shows how to generate 64 CIFAR10-like images and make a grid of them:
python3 -m megdiffusion.scripts.inference
Train from scratch
-
Take DDPM CIFAR10 for example:
python3 -m megdiffusion.scripts.train \ --flagfile ./megdiffusion/config/ddpm-cifar10.txt
-
[Optional] Overwrite arguments:
python3 -m megdiffusion.scripts.train \ --flagfile ./megdiffusion/config/ddpm-cifar10.txt \ --logdir ./path/to/logdir \ --batch_size=64 \ --save_step=100000 \ --parallel=True
Known issues:
- Training with single GPU & using gradient clipping will cause error in MegEngine 1.9.x version.
Development
python3 -m pip install -r requirements.txt
python3 -m pip install -v -e .
Develop this project with a new branch locally, remember to add necessary test codes.
If finished, submit Pull Request to the main
branch then just wait for review.
Acknowledgment
The following open-sourced projects was referenced here:
- hojonathanho/diffusion: The official Tensorflow implementation of DDPM.
- w86763777/pytorch-ddpm: Unofficial PyTorch implementation of Denoising Diffusion Probabilistic Models.
Thanks to people including @gaohuazuo, @xxr3376, @P2Oileen and other contributors for support in this project. The R&D platform and the resources required for the experiment are provided by MEGVII Inc. The deep learning framework used in this project is MegEngine -- a magic weapon.
Citations
@article{ho2020denoising,
title = {Denoising Diffusion Probabilistic Models},
author = {Jonathan Ho and Ajay Jain and Pieter Abbeel},
year = {2020},
eprint = {2006.11239},
archivePrefix = {arXiv},
primaryClass = {cs.LG}
}
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for megdiffusion-0.0.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 3a5f541d0921ada8ba660a495673cb73d39fb79e553b71d5e0f7c3263f1b41d3 |
|
MD5 | 77484281e83d572ab55b76f0dec91f5a |
|
BLAKE2b-256 | 9de481705c649d028b794991793e602ff14e85501d133fc55ec99b0bf1cdf1c0 |