alfafold3 pytorch
0.6.6
Implementación de Alphafold 3 en Pytorch
Puedes charlar con otros investigadores sobre este trabajo aquí
Revisión del artículo de Sergey
Guía ilustrada por Elana P. Simon
Charla de Max Jaderberg
Alex mantiene una bifurcación con soporte completo para Lightning + Hydra en este repositorio
Aquí se puede ver e interactuar con una visualización de las moléculas de vida utilizadas en el repositorio.
¡Joseph por contribuir con la codificación posicional relativa y la pérdida suave de LDDT!
Felipe por contribuir con los módulos de alineación rígida ponderada, coordenadas expresas en marco, error de alineación de cálculo y aumento aleatorio central.
Alex por solucionar varios problemas en los algoritmos transcritos.
Heng por señalar inconsistencias con el documento y solicitar soluciones.
Heng por encontrar un problema con los índices de átomos moleculares para la pérdida de distograma
Wei Lu por detectar algunos hiperparámetros erróneos
¡Alex por el script de preparación del conjunto de datos de PDB!
¡Milot por optimizar el script de agrupación de conjuntos de datos de PDB!
Alex, básicamente, por escribir todo el flujo gigantesco, desde el análisis del PDB hasta la molécula y las entradas atómicas para el entrenamiento.
¡Andrei por trabajar en el muestreo del conjunto de datos PDB ponderado!
Jimin por enviar una pequeña solución a un problema con las coordenadas que se pasan a WeightedRigidAlign
@xluo233 por contribuir con las medidas de confianza, la clasificación de penalizaciones por choque y la lógica de clasificación de muestra.
sj900 para integrar y probar WeightedPDBSampler dentro de PDBDataset y para agregar soporte inicial para MSA y análisis de plantillas.
@xluo233 nuevamente por contribuir con la lógica para calcular la puntuación de selección del modelo, así como el rasa no resuelto.
Fandi por descubrir algunas inconsistencias en el módulo de difusión atómica dilucidado con el suplemento
¡Paolo por proponer la hipótesis PDB neutral stable molecule !
¡Dhuvi por corregir un error relacionado con la asignación de ID de molécula de iones metálicos para Alphafold3Inputs !
¡Dhuvi por asumir la lógica para traducir Alphafold3Input a BioMolecule para guardarlo en mmCIF!
Tom (del canal Discord) por identificar una discrepancia entre los cálculos de distogramas y vectores unitarios de plantilla de esta base de código y los de OpenFold (y Andrei por ayudar a abordar el problema del distograma).
¡Kaihui por identificar un error en cómo se manejaban los átomos no estándar en los residuos de polímeros!
¡Andrei por encargarse de la interfaz frontend de gradio!
Patrick por jaxtyping, Florian por einx y, por supuesto, Alex por einops.
Soumith y la organización Pytorch por brindarme la oportunidad de abrir este trabajo.
$ pip install alphafold3-pytorch import torch
from alphafold3_pytorch import Alphafold3
from alphafold3_pytorch . utils . model_utils import exclusive_cumsum
alphafold3 = Alphafold3 (
dim_atom_inputs = 77 ,
dim_template_feats = 108
)
# mock inputs
seq_len = 16
molecule_atom_indices = torch . randint ( 0 , 2 , ( 2 , seq_len )). long ()
molecule_atom_lens = torch . full (( 2 , seq_len ), 2 ). long ()
atom_seq_len = molecule_atom_lens . sum ( dim = - 1 ). amax ()
atom_offsets = exclusive_cumsum ( molecule_atom_lens )
atom_inputs = torch . randn ( 2 , atom_seq_len , 77 )
atompair_inputs = torch . randn ( 2 , atom_seq_len , atom_seq_len , 5 )
additional_molecule_feats = torch . randint ( 0 , 2 , ( 2 , seq_len , 5 ))
additional_token_feats = torch . randn ( 2 , seq_len , 33 )
is_molecule_types = torch . randint ( 0 , 2 , ( 2 , seq_len , 5 )). bool ()
is_molecule_mod = torch . randint ( 0 , 2 , ( 2 , seq_len , 4 )). bool ()
molecule_ids = torch . randint ( 0 , 32 , ( 2 , seq_len ))
template_feats = torch . randn ( 2 , 2 , seq_len , seq_len , 108 )
template_mask = torch . ones (( 2 , 2 )). bool ()
msa = torch . randn ( 2 , 7 , seq_len , 32 )
msa_mask = torch . ones (( 2 , 7 )). bool ()
additional_msa_feats = torch . randn ( 2 , 7 , seq_len , 2 )
# required for training, but omitted on inference
atom_pos = torch . randn ( 2 , atom_seq_len , 3 )
distogram_atom_indices = molecule_atom_lens - 1
distance_labels = torch . randint ( 0 , 37 , ( 2 , seq_len , seq_len ))
resolved_labels = torch . randint ( 0 , 2 , ( 2 , atom_seq_len ))
# offset indices correctly
distogram_atom_indices += atom_offsets
molecule_atom_indices += atom_offsets
# train
loss = alphafold3 (
num_recycling_steps = 2 ,
atom_inputs = atom_inputs ,
atompair_inputs = atompair_inputs ,
molecule_ids = molecule_ids ,
molecule_atom_lens = molecule_atom_lens ,
additional_molecule_feats = additional_molecule_feats ,
additional_msa_feats = additional_msa_feats ,
additional_token_feats = additional_token_feats ,
is_molecule_types = is_molecule_types ,
is_molecule_mod = is_molecule_mod ,
msa = msa ,
msa_mask = msa_mask ,
templates = template_feats ,
template_mask = template_mask ,
atom_pos = atom_pos ,
distogram_atom_indices = distogram_atom_indices ,
molecule_atom_indices = molecule_atom_indices ,
distance_labels = distance_labels ,
resolved_labels = resolved_labels
)
loss . backward ()
# after much training ...
sampled_atom_pos = alphafold3 (
num_recycling_steps = 4 ,
num_sample_steps = 16 ,
atom_inputs = atom_inputs ,
atompair_inputs = atompair_inputs ,
molecule_ids = molecule_ids ,
molecule_atom_lens = molecule_atom_lens ,
additional_molecule_feats = additional_molecule_feats ,
additional_msa_feats = additional_msa_feats ,
additional_token_feats = additional_token_feats ,
is_molecule_types = is_molecule_types ,
is_molecule_mod = is_molecule_mod ,
msa = msa ,
msa_mask = msa_mask ,
templates = template_feats ,
template_mask = template_mask
)
sampled_atom_pos . shape # (2, , 3) Un ejemplo con manejo de entrada a nivel de molécula
import torch
from alphafold3_pytorch import Alphafold3 , Alphafold3Input
contrived_protein = 'AG'
mock_atompos = [
torch . randn ( 5 , 3 ), # alanine has 5 non-hydrogen atoms
torch . randn ( 4 , 3 ) # glycine has 4 non-hydrogen atoms
]
train_alphafold3_input = Alphafold3Input (
proteins = [ contrived_protein ],
atom_pos = mock_atompos
)
eval_alphafold3_input = Alphafold3Input (
proteins = [ contrived_protein ]
)
# training
alphafold3 = Alphafold3 (
dim_atom_inputs = 3 ,
dim_atompair_inputs = 5 ,
atoms_per_window = 27 ,
dim_template_feats = 108 ,
num_molecule_mods = 0 ,
confidence_head_kwargs = dict (
pairformer_depth = 1
),
template_embedder_kwargs = dict (
pairformer_stack_depth = 1
),
msa_module_kwargs = dict (
depth = 1
),
pairformer_stack = dict (
depth = 2
),
diffusion_module_kwargs = dict (
atom_encoder_depth = 1 ,
token_transformer_depth = 1 ,
atom_decoder_depth = 1 ,
)
)
loss = alphafold3 . forward_with_alphafold3_inputs ([ train_alphafold3_input ])
loss . backward ()
# sampling
alphafold3 . eval ()
sampled_atom_pos = alphafold3 . forward_with_alphafold3_inputs ( eval_alphafold3_input )
assert sampled_atom_pos . shape == ( 1 , ( 5 + 4 ), 3 ) Para adquirir el conjunto de datos PDB de AlphaFold 3, primero descargue todos los complejos de primer ensamblaje (y unidades asimétricas) en el Banco de datos de proteínas (PDB) y luego procéselos previamente con el script al que se hace referencia a continuación. El PDB se puede descargar desde RCSB: https://www.wwpdb.org/ftp/pdb-ftp-sites#rcsbpdb. Los dos scripts de Python siguientes (es decir, filter_pdb_{train,val,test}_mmcifs.py y cluster_pdb_{train,val,test}_mmcifs.py ) asumen que ha descargado el PDB en el formato de archivo mmCIF , colocando su primer ensamblado y archivos mmCIF de unidades asimétricas en data/pdb_data/unfiltered_assembly_mmcifs/ y data/pdb_data/unfiltered_asym_mmcifs/ , respectivamente.
Para lograr reproducibilidad, recomendamos descargar la PDB mediante instantáneas de AWS (por ejemplo, 20240101 ). Para hacerlo, consulte la documentación de AWS para configurar la AWS CLI localmente. Alternativamente, en el sitio web de RCSB, navegue hasta "Protocolos de descarga" y siga las instrucciones de descarga según su ubicación.
Por ejemplo, se pueden utilizar los siguientes comandos para descargar el PDB como dos colecciones de archivos mmCIF:
# For `assembly1` complexes, use the PDB's `20240101` AWS snapshot:
aws s3 sync s3://pdbsnapshots/20240101/pub/pdb/data/assemblies/mmCIF/divided/ ./data/pdb_data/unfiltered_assembly_mmcifs
# Or as a fallback, use rsync:
rsync -rlpt -v -z --delete --port=33444
rsync.rcsb.org::ftp_data/assemblies/mmCIF/divided/ ./data/pdb_data/unfiltered_assembly_mmcifs/
# For asymmetric unit complexes, also use the PDB's `20240101` AWS snapshot:
aws s3 sync s3://pdbsnapshots/20240101/pub/pdb/data/structures/divided/mmCIF/ ./data/pdb_data/unfiltered_asym_mmcifs
# Or as a fallback, use rsync:
rsync -rlpt -v -z --delete --port=33444
rsync.rcsb.org::ftp_data/structures/divided/mmCIF/ ./data/pdb_data/unfiltered_asym_mmcifs/ADVERTENCIA: La descarga del PDB puede ocupar hasta 700 GB de espacio.
NOTA: La PDB aloja todas las instantáneas de AWS disponibles aquí: https://pdbsnapshots.s3.us-west-2.amazonaws.com/index.html.
Después de la descarga, debería tener dos directorios formateados como este: https://files.rcsb.org/pub/pdb/data/assemblies/mmCIF/divided/ & https://files.rcsb.org/pub/pdb/data /estructuras/dividido/mmCIF/
00/
01/
02/
..
zz/Para estos directorios, descomprima todos los archivos:
find ./data/pdb_data/unfiltered_assembly_mmcifs/ -type f -name " *.gz " -exec gzip -d {} ;
find ./data/pdb_data/unfiltered_asym_mmcifs/ -type f -name " *.gz " -exec gzip -d {} ;Luego ejecute los comandos
wget -P ./data/ccd_data/ https://files.wwpdb.org/pub/pdb/data/monomers/components.cif.gz
wget -P ./data/ccd_data/ https://files.wwpdb.org/pub/pdb/data/component-models/complete/chem_comp_model.cif.gzdesde el directorio raíz del proyecto para descargar la última versión del Diccionario de componentes químicos (CCD) del PDB y sus modelos estructurales. Extraiga cada uno de estos archivos usando el siguiente comando:
find data/ccd_data/ -type f -name " *.gz " -exec gzip -d {} ; Luego ejecute lo siguiente con pdb_assembly_dir , pdb_asym_dir , ccd_dir y mmcif_output_dir reemplazados con las ubicaciones de sus copias locales del PDB del primer ensamblaje, PDB de la unidad asimétrica, CCD y el directorio de salida del conjunto de datos deseado (es decir, ./data/pdb_data/unfiltered_assembly_mmcifs/ , ./data/pdb_data/unfiltered_asym_mmcifs/ , ./data/ccd_data/ y ./data/pdb_data/{train,val,test}_mmcifs/ ).
python scripts/filter_pdb_train_mmcifs.py --mmcif_assembly_dir < pdb_assembly_dir > --mmcif_asym_dir < pdb_asym_dir > --ccd_dir < ccd_dir > --output_dir < mmcif_output_dir >
python scripts/filter_pdb_val_mmcifs.py --mmcif_assembly_dir < pdb_assembly_dir > --mmcif_asym_dir < pdb_asym_dir > --output_dir < mmcif_output_dir >
python scripts/filter_pdb_test_mmcifs.py --mmcif_assembly_dir < pdb_assembly_dir > --mmcif_asym_dir < pdb_asym_dir > --output_dir < mmcif_output_dir > Consulte los scripts para obtener más opciones. Cada mmCIF del primer ensamblado que pase exitosamente todos los pasos de procesamiento se escribirá en mmcif_output_dir dentro de un subdirectorio nombrado de acuerdo con el segundo y tercer carácter ID de PDB del mmCIF (por ejemplo, 5c ).
A continuación, ejecute lo siguiente con mmcif_dir y {train,val,test}_clustering_output_dir reemplazados, respectivamente, con su directorio de salida local creado usando el script de filtrado de conjuntos de datos anterior y con los directorios de salida de agrupación que desee (es decir, ./data/pdb_data/{train,val,test}_mmcifs/ y ./data/pdb_data/data_caches/{train,val,test}_clusterings/ ):
python scripts/cluster_pdb_train_mmcifs.py --mmcif_dir < mmcif_dir > --output_dir < train_clustering_output_dir > --clustering_filtered_pdb_dataset
python scripts/cluster_pdb_val_mmcifs.py --mmcif_dir < mmcif_dir > --reference_clustering_dir < train_clustering_output_dir > --output_dir < val_clustering_output_dir > --clustering_filtered_pdb_dataset
python scripts/cluster_pdb_test_mmcifs.py --mmcif_dir < mmcif_dir > --reference_1_clustering_dir < train_clustering_output_dir > --reference_2_clustering_dir < val_clustering_output_dir > --output_dir < test_clustering_output_dir > --clustering_filtered_pdb_dataset Nota : Se recomienda el indicador --clustering_filtered_pdb_dataset al agrupar el conjunto de datos PDB filtrado según lo seleccionado mediante los scripts anteriores, ya que este indicador permitirá tiempos de ejecución más rápidos en este contexto (ya que el filtrado deja los ID de residuos de cada cadena basados en 1). Sin embargo, este indicador no se debe proporcionar al agrupar otros conjuntos de datos (es decir, que no sean PDB) de archivos mmCIF. De lo contrario, es posible que la agrupación de interfaces se realice incorrectamente, ya que es posible que los archivos mmCIF de estos conjuntos de datos no utilicen una indexación estricta de residuos basada en 1 para cada cadena.
Nota : En su lugar, se pueden descargar archivos mmCIF ( train / val / test ) preprocesados (es decir, filtrados) (~25 GB, que comprenden 148 k complejos) y archivos de agrupación en cadena/interfaz ( train / val / test ) (~3 GB) para 20240101 de PDB. Instantánea de AWS a través de una carpeta compartida de OneDrive. Cada uno de estos archivos tar.gz debe descomprimirse dentro del directorio data/pdb_data/ , por ejemplo, a través de tar -xzf data_caches.tar.gz -C data/pdb_data/ . También se pueden descargar y preparar datos de destilación de PDB utilizando como referencia el script scripts/distillation_data_download.sh . Una vez descargado, se puede ejecutar scripts/reduce_uniprot_predictions_to_pdb.py para filtrar este conjunto de datos solo a ejemplos asociados con al menos una entrada de PDB. Además, para mayor comodidad, ya se ha descargado y extraído una asignación de los ID de acceso de UniProt a los ID de PDB para la capacitación sobre datos de destilación de PDB como data/afdb_data/data_caches/uniprot_to_pdb_id_mapping.dat .
En la raíz del proyecto, ejecute
$ sh ./contribute.sh Luego, agregue su módulo a alphafold3_pytorch/alphafold3.py , agregue sus pruebas a tests/test_af3.py y envíe una solicitud de extracción. Puede ejecutar las pruebas localmente con
$ pytest tests/ El Dockerfile incluido contiene las dependencias necesarias para ejecutar el paquete y entrenar/inferir usando PyTorch con GPU.
La imagen base predeterminada es pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime e instala la última versión de este paquete desde la rama main de GitHub.
# # Build Docker Container
docker build -t af3 .Alternativamente, use argumentos de compilación para reconstruir la imagen con diferentes versiones de software:
PYTORCH_TAG : cambia la imagen base y, por lo tanto, compila con diferentes versiones de PyTorch, CUDA y/o cuDNN.GIT_TAG : cambia la etiqueta de este repositorio para clonar e instalar el paquete.Por ejemplo:
# # Use build argument to change versions
docker build --build-arg " PYTORCH_TAG=2.2.1-cuda12.1-cudnn8-devel " --build-arg " GIT_TAG=0.1.15 " -t af3 .Luego, ejecute el contenedor con GPU y monte un volumen local (para entrenamiento) usando el siguiente comando:
# # Run Container
docker run -v .:/data --gpus all -it af3 @article { Abramson2024-fj ,
title = " Accurate structure prediction of biomolecular interactions with
{AlphaFold} 3 " ,
author = " Abramson, Josh and Adler, Jonas and Dunger, Jack and Evans,
Richard and Green, Tim and Pritzel, Alexander and Ronneberger,
Olaf and Willmore, Lindsay and Ballard, Andrew J and Bambrick,
Joshua and Bodenstein, Sebastian W and Evans, David A and Hung,
Chia-Chun and O'Neill, Michael and Reiman, David and
Tunyasuvunakool, Kathryn and Wu, Zachary and {v Z}emgulyt{.e},
Akvil{.e} and Arvaniti, Eirini and Beattie, Charles and
Bertolli, Ottavia and Bridgland, Alex and Cherepanov, Alexey and
Congreve, Miles and Cowen-Rivers, Alexander I and Cowie, Andrew
and Figurnov, Michael and Fuchs, Fabian B and Gladman, Hannah and
Jain, Rishub and Khan, Yousuf A and Low, Caroline M R and Perlin,
Kuba and Potapenko, Anna and Savy, Pascal and Singh, Sukhdeep and
Stecula, Adrian and Thillaisundaram, Ashok and Tong, Catherine
and Yakneen, Sergei and Zhong, Ellen D and Zielinski, Michal and
{v Z}{'i}dek, Augustin and Bapst, Victor and Kohli, Pushmeet
and Jaderberg, Max and Hassabis, Demis and Jumper, John M " ,
journal = " Nature " ,
month = " May " ,
year = 2024
} @inproceedings { Darcet2023VisionTN ,
title = { Vision Transformers Need Registers } ,
author = { Timoth'ee Darcet and Maxime Oquab and Julien Mairal and Piotr Bojanowski } ,
year = { 2023 } ,
url = { https://api.semanticscholar.org/CorpusID:263134283 }
} @article { Arora2024SimpleLA ,
title = { Simple linear attention language models balance the recall-throughput tradeoff } ,
author = { Simran Arora and Sabri Eyuboglu and Michael Zhang and Aman Timalsina and Silas Alberti and Dylan Zinsley and James Zou and Atri Rudra and Christopher R'e } ,
journal = { ArXiv } ,
year = { 2024 } ,
volume = { abs/2402.18668 } ,
url = { https://api.semanticscholar.org/CorpusID:268063190 }
} @article { Puny2021FrameAF ,
title = { Frame Averaging for Invariant and Equivariant Network Design } ,
author = { Omri Puny and Matan Atzmon and Heli Ben-Hamu and Edward James Smith and Ishan Misra and Aditya Grover and Yaron Lipman } ,
journal = { ArXiv } ,
year = { 2021 } ,
volume = { abs/2110.03336 } ,
url = { https://api.semanticscholar.org/CorpusID:238419638 }
} @article { Duval2023FAENetFA ,
title = { FAENet: Frame Averaging Equivariant GNN for Materials Modeling } ,
author = { Alexandre Duval and Victor Schmidt and Alex Hernandez Garcia and Santiago Miret and Fragkiskos D. Malliaros and Yoshua Bengio and David Rolnick } ,
journal = { ArXiv } ,
year = { 2023 } ,
volume = { abs/2305.05577 } ,
url = { https://api.semanticscholar.org/CorpusID:258564608 }
} @article { Wang2022DeepNetST ,
title = { DeepNet: Scaling Transformers to 1, 000 Layers } ,
author = { Hongyu Wang and Shuming Ma and Li Dong and Shaohan Huang and Dongdong Zhang and Furu Wei } ,
journal = { ArXiv } ,
year = { 2022 } ,
volume = { abs/2203.00555 } ,
url = { https://api.semanticscholar.org/CorpusID:247187905 }
} @inproceedings { Ainslie2023CoLT5FL ,
title = { CoLT5: Faster Long-Range Transformers with Conditional Computation } ,
author = { Joshua Ainslie and Tao Lei and Michiel de Jong and Santiago Ontan'on and Siddhartha Brahma and Yury Zemlyanskiy and David Uthus and Mandy Guo and James Lee-Thorp and Yi Tay and Yun-Hsuan Sung and Sumit Sanghai } ,
year = { 2023 }
} @article { Ash2019OnTD ,
title = { On the Difficulty of Warm-Starting Neural Network Training } ,
author = { Jordan T. Ash and Ryan P. Adams } ,
journal = { ArXiv } ,
year = { 2019 } ,
volume = { abs/1910.08475 } ,
url = { https://api.semanticscholar.org/CorpusID:204788802 }
} @ARTICLE { Heinzinger2023.07.23.550085 ,
author = { Michael Heinzinger and Konstantin Weissenow and Joaquin Gomez Sanchez and Adrian Henkel and Martin Steinegger and Burkhard Rost } ,
title = { ProstT5: Bilingual Language Model for Protein Sequence and Structure } ,
year = { 2023 } ,
doi = { 10.1101/2023.07.23.550085 } ,
journal = { bioRxiv }
} @article { Lin2022.07.20.500902 ,
author = { Lin, Zeming and Akin, Halil and Rao, Roshan and Hie, Brian and Zhu, Zhongkai and Lu, Wenting and Santos Costa, Allan dos and Fazel-Zarandi, Maryam and Sercu, Tom and Candido, Sal and Rives, Alexander } ,
title = { Language models of protein sequences at the scale of evolution enable accurate structure prediction } ,
elocation-id = { 2022.07.20.500902 } ,
year = { 2022 } ,
doi = { 10.1101/2022.07.20.500902 } ,
publisher = { Cold Spring Harbor Laboratory } ,
URL = { https://www.biorxiv.org/content/early/2022/07/21/2022.07.20.500902 } ,
eprint = { https://www.biorxiv.org/content/early/2022/07/21/2022.07.20.500902.full.pdf } ,
journal = { bioRxiv }
} @article { Li2024SwitchEA ,
title = { Switch EMA: A Free Lunch for Better Flatness and Sharpness } ,
author = { Siyuan Li and Zicheng Liu and Juanxi Tian and Ge Wang and Zedong Wang and Weiyang Jin and Di Wu and Cheng Tan and Tao Lin and Yang Liu and Baigui Sun and Stan Z. Li } ,
journal = { ArXiv } ,
year = { 2024 } ,
volume = { abs/2402.09240 } ,
url = { https://api.semanticscholar.org/CorpusID:267657558 }
} @article { Nguyen2023MitigatingOI ,
title = { Mitigating Over-smoothing in Transformers via Regularized Nonlocal Functionals } ,
author = { Tam Nguyen and Tan M. Nguyen and Richard G. Baraniuk } ,
journal = { ArXiv } ,
year = { 2023 } ,
volume = { abs/2312.00751 } ,
url = { https://api.semanticscholar.org/CorpusID:264300597 }
} @inproceedings { Zhou2024ValueRL ,
title = { Value Residual Learning For Alleviating Attention Concentration In Transformers } ,
author = { Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan } ,
year = { 2024 } ,
url = { https://api.semanticscholar.org/CorpusID:273532030 }
}
