Denoising Diffusion Probabilistic Model for Proteins
Implementation of Denoising Diffusion Probabilistic Model in Pytorch. It is a new approach to generative modeling that may have the potential to rival GANs. It uses denoising score matching to estimate the gradient of the data distribution, followed by Langevin sampling to sample from the true distribution. This implementation was transcribed from the official Tensorflow version here.
This specific repository will be using a heavily modifying version of the U-net for learning on protein structure, with eventual conditioning from MSA Transformers attention heads.
** at around 40k iterations **
Install
$ pip install ddpm-proteins
Training
We are using weights & biases for experimental tracking
First you need to login
$ wandb login
Then
$ python train.py
Edit train.py to whatever for your research desires
Todo
- condition on mask
- condition on MSA transformers (with caching of tensors in specified directory by protein id)
- reach for size 384
- all-attention network with uformer https://arxiv.org/abs/2106.03106 (with 1d + 2d conv kernels)
- add all improvements from https://arxiv.org/abs/2105.05233 and https://cascaded-diffusion.github.io/
Usage
import torch
from ddpm_proteins import Unet, GaussianDiffusion
model = Unet(
dim = 64,
dim_mults = (1, 2, 4, 8)
)
diffusion = GaussianDiffusion(
model,
image_size = 128,
timesteps = 1000, # number of steps
loss_type = 'l1' # L1 or L2
)
training_images = torch.randn(8, 3, 128, 128)
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)
Or, if you simply want to pass in a folder name and the desired image dimensions, you can use the Trainer class to easily train a model.
from ddpm_proteins import Unet, GaussianDiffusion, Trainer
model = Unet(
dim = 64,
dim_mults = (1, 2, 4, 8)
).cuda()
diffusion = GaussianDiffusion(
model,
image_size = 128,
timesteps = 1000, # number of steps
loss_type = 'l1' # L1 or L2
).cuda()
trainer = Trainer(
diffusion,
'path/to/your/images',
train_batch_size = 32,
train_lr = 2e-5,
train_num_steps = 700000, # total training steps
gradient_accumulate_every = 2, # gradient accumulation steps
ema_decay = 0.995, # exponential moving average decay
fp16 = True # turn on mixed precision training with apex
)
trainer.train()
Samples and model checkpoints will be logged to ./results periodically
Citations
@misc{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}
}
@inproceedings{anonymous2021improved,
title = {Improved Denoising Diffusion Probabilistic Models},
author = {Anonymous},
booktitle = {Submitted to International Conference on Learning Representations},
year = {2021},
url = {https://openreview.net/forum?id=-NEXDKk8gZ},
note = {under review}
}
@article{Rao2021.02.12.430858,
author = {Rao, Roshan and Liu, Jason and Verkuil, Robert and Meier, Joshua and Canny, John F. and Abbeel, Pieter and Sercu, Tom and Rives, Alexander},
title = {MSA Transformer},
year = {2021},
publisher = {Cold Spring Harbor Laboratory},
URL = {https://www.biorxiv.org/content/early/2021/02/13/2021.02.12.430858},
journal = {bioRxiv}
}
581 Dec 30, 2022
68 Dec 26, 2022
2.9k Jan 4, 2023
155 Dec 26, 2022
71 Dec 1, 2022
36 Nov 29, 2022
2 Mar 7, 2022
3k Jan 8, 2023
46 Nov 24, 2022
1.5k Dec 29, 2022
24 Dec 14, 2022
1 Dec 18, 2021
310 Nov 07, 2022
4 Jul 12, 2022
113 Nov 29, 2022
357 Jan 06, 2023
17k Jan 02, 2023
12 Jan 13, 2022
53 Sep 17, 2022
920 Jan 04, 2023
215 Dec 17, 2022
26 Dec 27, 2022
1 Nov 07, 2021
322 Dec 23, 2022
20 Oct 18, 2021
237 Jan 05, 2023
1.4k Dec 30, 2022
152 Dec 22, 2022