OpenFold / README.md

Upload folder using huggingface_hub

d149fb1 verified 1 day ago

8.7 kB

license: apache-2.0
language:
  - en
  - zh
tags:
  - OneScience
  - life-science
  - protein-structure-prediction
frameworks: PyTorch

OpenFold

Model Introduction

OpenFold is an open-source framework for predicting three-dimensional protein structures from protein sequences, multiple sequence alignments (MSAs), and template information. As a trainable reproduction of AlphaFold2, its core network architecture includes Evoformer, Structure Module, template modules, MSA modules, and other components, supporting a complete structural prediction experiment workflow that includes training, inference, and weight conversion.

Paper: OpenFold: Retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization https://www.biorxiv.org/content/10.1101/2022.11.20.517210v2

Model Description

OpenFold is based on the Evoformer + Structure Module architecture and is trained with PDB and UniRef data. It is designed for protein three-dimensional structure prediction, model fine-tuning, and research on architectural improvements. The core network contains 48 Evoformer Blocks and 8 Structure Module layers. Through triangle attention and invariant point attention mechanisms, it enables bidirectional fusion of evolutionary information and spatial geometry. Compared with the original AlphaFold2, OpenFold adds memory-efficient attention, mixed-precision training support, and seamless weight conversion capabilities, significantly reducing GPU memory usage and fixing trainability issues.

Applicable Scenarios

Scenario	Description
Local OpenFold invocation	Users can call OpenFold training or inference through the scripts in this repository.
Reuse of the OneScience base	Continue using shared OneScience modules such as datapipes, utils, loss, and np.
ModelScope/OneCode release	Expose only scripts, models, configurations, and weight directories to reduce the release package size.
Weight conversion and data preparation	Use the tools under `scripts/` to download databases, preprocess alignments, and convert weights.

Usage

1. Using OneCode

You can try intelligent one-click AI4S programming through the OneCode online environment:

Try intelligent one-click AI4S programming

2. Manual Installation and Usage

Hardware Requirements

Running on a GPU or DCU is recommended.
A CPU can be used for import checks and small-configuration connectivity validation, but full training and inference will be slow.
DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version that matches the current cluster.

3. Quick Start

Download the Model Package

modelscope download --model OneScience/OpenFold --local_dir ./OpenFold
cd OpenFold

Install the Runtime Environment

DCU Environment

# Activate DTK and CONDA first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# uv installation is supported
pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Quick Verification

PYTHONPATH=$(pwd)/model  python -c "from openfold.model import AlphaFold; from config.config import model_config; import onescience; print('openfold wrapper ok')"

If dependencies such as torch, ml_collections, or deepspeed are reported as missing, install the OneScience-recommended runtime environment first, or install the corresponding dependencies for your DCU/GPU version.

Example Data

The data/ directory is used to store the example inputs required by OpenFold and user-provided data. The current repository includes a monomer inference example at data/demo/monomer, making it easy to quickly check the inference script, model import, and alignment loading workflow.

Directory structure:

Path	Description
`data/demo/monomer/fasta_dir/6kwc.fasta`	Example protein sequence. The FASTA header is `6KWC_1`.
`data/demo/monomer/alignments/6KWC_1/`	Precomputed MSA/template search results corresponding to the FASTA header.
`data/demo/monomer/alignments/6KWC_1/bfd_uniref_hits.a3m`	BFD/UniRef alignment results.
`data/demo/monomer/alignments/6KWC_1/uniref90_hits.sto`	UniRef90 alignment results.
`data/demo/monomer/alignments/6KWC_1/mgnify_hits.sto`	MGnify alignment results.
`data/demo/monomer/alignments/6KWC_1/hhsearch_output.hhr`	HHsearch template search results.
`data/demo/monomer/inference.sh`	Monomer inference example script using the FASTA and precomputed alignments above.

The example data does not include OpenFold weights or the PDB mmCIF template library. Before running the demo, prepare:

weight/openfold.pt, or specify the weights with CHECKPOINT_PATH=/path/to/openfold.pt.
The PDB mmCIF template directory. By default, data/databases/pdb_mmcif/mmcif_files is read; it can also be specified with TEMPLATE_MMCIF_DIR=/path/to/mmcif_files.

Run the example:

bash data/demo/monomer/inference.sh

To use another device, set MODEL_DEVICE:

MODEL_DEVICE=cuda:0 bash data/demo/monomer/inference.sh

Inference Example

python scripts/inference.py /path/to/fasta_dir /path/to/template_mmcif_dir \
  --output_dir ./outputs/inference \
  --config_preset model_1_ptm \
  --model_device cuda:0 \
  --use_precomputed_alignments /path/to/alignments \
  --openfold_checkpoint_path ./weight/openfold.pt

If using AlphaFold JAX parameters:

python scripts/inference.py /path/to/fasta_dir /path/to/template_mmcif_dir \
  --output_dir ./outputs/inference \
  --config_preset model_1_ptm \
  --model_device cuda:0 \
  --use_precomputed_alignments /path/to/alignments \
  --jax_param_path ./weight/params_model_1_ptm.npz

Training Example

python scripts/train.py \
  /path/to/train_mmcif \
  /path/to/train_alignments \
  /path/to/template_mmcif \
  ./outputs/train \
  2021-10-10 \
  --config_preset initial_training \
  --max_epochs 1 \
  --train_epoch_len 1 \
  --gpus 1

Multi-GPU training can be launched with torchrun; adjust the specific parameters according to your runtime environment.

Data and Weight Downloads

OpenFold weights:

bash scripts/download_openfold_params.sh ./weights

Hugging Face weights:

bash scripts/download_openfold_params_huggingface.sh ./weights

PDB mmCIF template library:

bash scripts/download_pdb_mmcif.sh /path/to/database_dir

Complete AlphaFold/OpenFold databases:

bash scripts/download_alphafold_dbs.sh /path/to/database_dir full_dbs

Official OneScience Information

Platform	OneScience Main Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citations and License

Original OpenFold paper: OpenFold: Retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization.
Original AlphaFold2 paper: Highly accurate protein structure prediction with AlphaFold.
If you use OpenFold's multimer prediction functionality, you should also cite the original AlphaFold-Multimer paper: Protein complex prediction with AlphaFold-Multimer.
If you use OpenProteinSet training data, you should also cite: OpenProteinSet: Training data for structural biology at scale.
OpenFold-related source code uses the Apache License 2.0. See LICENSE in the repository root for details. The terms of use for model weights and data should follow the instructions provided by the corresponding publishers.
The OpenProteinSet dataset uses the CC BY 4.0 License. When using this dataset, you should acknowledge the data source and cite the OpenProteinSet paper as required by the dataset page.
If you use OpenFold in research, it is recommended to cite the original OpenFold and AlphaFold2 papers as well as relevant OneScience project information. When using multimer functionality, also cite AlphaFold-Multimer. When using OpenProteinSet or other downstream data resources, also include citations for the corresponding datasets, databases, and original papers.