README.md 6.48 KB
Newer Older
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
1
2
3
4
5
# OpenFold

A faithful PyTorch reproduction of DeepMind's 
[AlphaFold 2](https://github.com/deepmind/alphafold).

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
6
7
8
9
10
11
## Features

OpenFold carefully reproduces (almost) all of the features of the original open
source inference code. The sole exception is model ensembling, which fared
poorly in DeepMind's own ablation testing and is being phased out in future
DeepMind experiments. It is omitted here for the sake of reducing clutter. In 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
12
cases where the *Nature* paper differs from the source, we always defer to the 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
13
14
15
latter. 

OpenFold is built to support inference with AlphaFold's original JAX weights.
16
17
Try it out with our [Colab notebook](https://colab.research.google.com/github/aqlaboratory/openfold/blob/main/notebooks/OpenFold.ipynb)
(not yet visible from Colab because the repo is still private).
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
18
19
20

Unlike DeepMind's public code, OpenFold is also trainable. It can be trained 
with or without [DeepSpeed](https://github.com/microsoft/deepspeed) and with 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
21
mixed precision. `bfloat16` training is not currently supported, but will be 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
22
in the future.
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
23

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
24
## Installation (Linux)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
25

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
26
Python dependencies available through `pip` are provided in `requirements.txt`. 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
27
28
OpenFold depends on `openmm==7.5.1` and `pdbfixer`, which are only available 
via `conda`. For producing sequence alignments, you'll also need `jackhmmer`, 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
29
`kalign`, and the [HH-suite](https://github.com/soedinglab/hh-suite) installed 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
30
on your system. Finally, some download scripts require `aria2c`.
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
31

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
32
For convenience, we provide a script that installs Miniconda locally, creates a 
33
34
`conda` virtual environment, installs all Python dependencies, and downloads
useful resources (including DeepMind's pretrained parameters). Run:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
35
36

```bash
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
37
38
39
scripts/install_third_party_dependencies.sh
```

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
40
To activate the environment, run:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
41
42

```bash
sft-managed's avatar
sft-managed committed
43
source scripts/activate_conda_env.sh
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
44
45
```

46
To deactivate it, run:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
47
48

```bash
sft-managed's avatar
sft-managed committed
49
source scripts/deactivate_conda_env.sh
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
50
51
```

52
53
54
55
56
57
To install the HH-suite to `/usr/bin`, run

```bash
# scripts/install_hh_suite.sh
```

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
58
## Usage
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
59

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
60
61
62
63
64
65
To download the genetic databases used by AlphaFold/OpenFold, run:

```bash
scripts/download_all_data.sh data/
```

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
66
### Inference
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
67

sft-managed's avatar
sft-managed committed
68
To run inference on a sequence `target.fasta` (e.g., `wget https://www.rcsb.org/fasta/entry/4DSN`) using a set of DeepMind's pretrained parameters, run e.g.
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
69

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
70
```bash
71
python3 run_pretrained_openfold.py \
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
72
    target.fasta \
73
74
75
76
77
78
79
    data/uniref90/uniref90.fasta \
    data/mgnify/mgy_clusters_2018_12.fa \
    data/pdb70/pdb70 \
    data/pdb_mmcif/mmcif_files/ \
    data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
    --output_dir ./ \
    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
sft-managed's avatar
sft-managed committed
80
81
82
83
84
    --device cuda:1 \
    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
    --hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
    --hhsearch_binary_path lib/conda/envs/openfold_venv/bin/hhsearch \
    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
85
```
86

sft-managed's avatar
sft-managed committed
87
where `data` is the same directory as in the previous step. If `jackhmmer`, `hhblits`, `hhsearch` and `kalign` are available at the default path of `/usr/bin`, their `binary_path` command-line arguments can be dropped. 
88

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
89
### Training
90

sft-managed's avatar
sft-managed committed
91
92
93
94
95
After activating the OpenFold environment with `source scripts/activate_conda_env.sh`, install OpenFold by running
```bash
python setup.py install
```
To train the model, you will first need to precompute protein alignments. Create `mmcif_dir/` and download `.cif` files from the PDB (e.g., `wget https://files.rcsb.org/download/4DSN.cif`). Then run:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
96
97
98
99
100
101
102
103
104

```bash
python3 scripts/precompute_alignments.py mmcif_dir/ alignment_dir/ \
    data/uniref90/uniref90.fasta \
    data/mgnify/mgy_clusters_2018_12.fa \
    data/pdb70/pdb70 \
    data/pdb_mmcif/mmcif_files/ \
    data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
    --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
sft-managed's avatar
sft-managed committed
105
106
107
108
109
    --cpus 16 \
    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
    --hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
    --hhsearch_binary_path lib/conda/envs/openfold_venv/bin/hhsearch \
    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
110
```
sft-managed's avatar
sft-managed committed
111
As noted before, you can skip the `binary_path` arguments if these binaries are at `/usr/bin`.
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
112
113
Expect this step to take a very long time, even for small numbers of proteins.

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
114
Next, generate a cache of certain datapoints in the mmCIF files:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
115
116

```bash
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
117
118
119
120
python3 scripts/generate_mmcif_cache.py \
    mmcif_dir/ \
    mmcif_cache.json \
    --no_workers 16
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
```

This cache is used to minimize the number of mmCIF parses performed during 
training-time data preprocessing. Finally, call the training script:

```bash
python3 train_openfold.py mmcif_dir/ alignment_dir/ template_mmcif_dir/ \
    2021-10-10 \ 
    --template_release_dates_cache_path mmcif_cache.json \ 
    --precision 16 \
    --gpus 8 --replace_sampler_ddp=True \
    --accelerator ddp \ 
    --seed 42 \ # in multi-gpu settings, the seed must be specified
    --deepspeed_config_path deepspeed_config.json
```

where `--template_release_dates_cache_path` is a path to the `.json` file
generated in the previous step. A suitable DeepSpeed configuration file can be 
generated with `scripts/build_deepspeed_config.py`. The training script is 
written with [PyTorch Lightning](https://github.com/PyTorchLightning/pytorch-lightning) 
and supports the full range of training options that entails, including 
multi-node distributed training. For more information, consult PyTorch 
Lightning documentation and the `--help` flag of the training script.

145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
## Testing

To run unit tests, use

```bash
scripts/run_unit_tests.sh
```

The script is a thin wrapper around Python's `unittest` suite, and recognizes
`unittest` commands. E.g., to run a specific test verbosely:

```bash
scripts/run_unit_tests.sh -v tests.test_model
```

Certain tests require that AlphaFold be installed in the same Python
environment. These run components of AlphaFold and OpenFold side by side and
ensure that output activations are adequately similar. For most modules, we
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
163
target a maximum difference of `1e-4`.
164

165
166
167
168
169
170
171
## Copyright notice

While AlphaFold's and, by extension, OpenFold's source code is licensed under
the permissive Apache Licence, Version 2.0, DeepMind's pretrained parameters 
remain under the more restrictive CC BY-NC 4.0 license, a copy of which is 
downloaded to `openfold/resources/params` by the installation script. They are
thereby made unavailable for commercial use.
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
172
173
174
175

## Contributing

If you encounter problems using OpenFold, feel free to create an issue!