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
7 Feb 20, 2022
1 Nov 09, 2021
19 Jan 07, 2023
192 Dec 26, 2022
150 Nov 14, 2022
3 Jun 23, 2022
68 Dec 23, 2022
61 Jan 03, 2023
8 Nov 29, 2022
50 Sep 26, 2022
7 Sep 12, 2022
54 Oct 09, 2022
130 Jan 06, 2023
15 Dec 17, 2022
64 Nov 21, 2022
107 Dec 27, 2022
0 Dec 13, 2021
172 Dec 29, 2022
470 Dec 30, 2022
111 Nov 24, 2022