ESMFold Structure Prediction
After installing with the [esmfold] option, you can use the ESMFold structure prediction model as follows:
import torch
import esm
model = esm.pretrained.esmfold_v1()
model = model.eval().cuda()
# Optionally, uncomment to set a chunk size for axial attention. This can help reduce memory.
# Lower sizes will have lower memory requirements at the cost of increased speed.
# model.set_chunk_size(128)
sequence = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
# Multimer prediction can be done with chains separated by ':'
with torch.no_grad():
output = model.infer_pdb(sequence)
with open("result.pdb", "w") as f:
f.write(output)
import biotite.structure.io as bsio
struct = bsio.load_structure("result.pdb", extra_fields=["b_factor"])
print(struct.b_factor.mean()) # this will be the pLDDT
# 88.3
Besides esm.pretrained.esmfold_v1() which is the best performing model we recommend using, we also provide esm.pretrained.esmfold_v0() which was used for the experiments in Lin et al. 2022.
We also provide a command line interface (esm-fold) that efficiently predicts structures in bulk from a FASTA file using ESMFold:
usage: esm-fold [-h] -i FASTA -o PDB [--num-recycles NUM_RECYCLES]
[--max-tokens-per-batch MAX_TOKENS_PER_BATCH]
[--chunk-size CHUNK_SIZE] [--cpu-only] [--cpu-offload]
optional arguments:
-h, --help show this help message and exit
-i FASTA, --fasta FASTA
Path to input FASTA file
-o PDB, --pdb PDB Path to output PDB directory
--num-recycles NUM_RECYCLES
Number of recycles to run. Defaults to number used in
training (4).
--max-tokens-per-batch MAX_TOKENS_PER_BATCH
Maximum number of tokens per gpu forward-pass. This
will group shorter sequences together for batched
prediction. Lowering this can help with out of memory
issues, if these occur on short sequences.
--chunk-size CHUNK_SIZE
Chunks axial attention computation to reduce memory
usage from O(L^2) to O(L). Equivalent to running a for
loop over chunks of of each dimension. Lower values
will result in lower memory usage at the cost of
speed. Recommended values: 128, 64, 32. Default: None.
--cpu-only CPU only
--cpu-offload Enable CPU offloading
The command will make one prediction for every sequence in the fasta file. Multimers can be predicted and should be entered in the fasta file as a single sequence, with chains seprated by a ":" character.
By default, predictions will be batched together so that shorter sequences are predicted simultaneously. This can be disabled by setting --max-tokens-per-batch=0. Batching can significantly improve prediction speed on shorter sequences.
The --cpu-offload flag can be useful for making predictions on longer sequences. It will attempt to offload some parameters to the CPU RAM, rather than storing on GPU.
Finally, the ablation experiments for LMs of varying sizes Lin et al. 2022 table S1 are released as esm.pretrained.esmfold_structure_module_only_*(). We don't recommend using these models for structure prediction.
Address: 800 Dong Chuan RD. Minhang District, Shanghai, China SJTU-Yale Joint Center for Biostatistics, SJTU