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
367 Dec 26, 2022
11 Jan 01, 2023
32 Jan 06, 2023
8 Sep 27, 2022
8 Dec 13, 2022
37 Dec 28, 2022
51 Nov 22, 2022
297 Dec 12, 2022
19 Dec 11, 2022
225 Dec 28, 2022
22 Oct 20, 2022
1.5k Dec 31, 2022
10 Nov 02, 2022
297 Dec 17, 2022
268 Jan 09, 2023
33 Dec 08, 2022
132 Nov 21, 2022
134 Jan 01, 2023
1 Dec 28, 2021
103 Dec 28, 2022