DOCSEARCHbyAlgolia
Submit Docs

RoseTTAFold2NA (RF2NA)

repository·main·Indexed 18 days ago

https://github.com/uw-ipd/rosettafold2na

A deep learning framework for predicting the structures of complexes involving proteins and nucleic acids (RNA and DNA). The repository includes the RF2NA prediction pipeline, instructions for downloading required sequence and structure databases (UniRef30, BFD, Rfam, etc.), and an implementation of the SE(3)-Transformer, an equivariant Graph Neural Network built with PyTorch and DGL for processing 3D points and graphs.

Tokens
4.3K
Snippets
11
Records
19
Agent score
14%

What's inside RoseTTAFold2NA

  1. Overview of SE(3)-Transformers For PyTorch

    main
    This repository provides scripts and recipes for training the SE(3)-Transformer model to achieve state-of-the-art accuracy. The content is maintained and tested by NVIDIA. It includes details on model architecture, default configurations, feature support (including mixed precision training), and performance benchmarking for both training and inference on NVIDIA hardware (e.g., DGX A100, DGX-1).

  2. What is the SE(3)-Transformer model?

    main

    The SE(3)-Transformer is a Graph Neural Network (GNN) that uses a variant of self-attention for processing 3D points and graphs. It is equivariant under continuous 3D roto-translations (SE(3)), meaning that if the input graphs or point sets rotate or undergo proper rigid transformations, the model outputs either remain invariant or transform predictably with the input.

    Key characteristics include:

    • Architecture: Consists of stacked layers of equivariant graph self-attention and equivariant normalization, followed by a Tensor Field Network (TFN) convolution to obtain invariant features. It typically functions as a graph-to-scalar network using graph pooling (mean or max).
    • Implementation: Built using PyTorch and the Deep Graph Library (DGL), utilizing the e3nn library for computing spherical harmonics and Clebsch–Gordan coefficients.
    • Capabilities: Supports edge features of any degree (1D, 3D, 5D, etc.) and is optimized for high throughput and low memory consumption.

  3. How to enable TF32 on NVIDIA Ampere GPUs

    main

    TensorFloat-32 (TF32) is a math mode for NVIDIA A100 GPUs that provides up to 10x speedups for matrix operations compared to single-precision (FP32) on older architectures. It is more robust than FP16 for models requiring high dynamic range.

    TF32 is enabled by default in the NVIDIA Ampere GPU architecture within this implementation.

  4. Understand SE(3)-Transformer feature terminology

    main

    When working with the model, the following terms define the data structures and mathematical properties:

    • Degree (type): A positive integer describing how a feature transforms under 3D rotation. A feature of degree $d$ has a dimensionality of $2d+1$.
      • Degree 0: 1D scalars (rotation invariant).
      • Degree 1: 3D vectors (rotate via 3D rotation matrices).
      • Degree 2: 5D vectors (rotate via Wigner-D matrices).
    • Fiber: A representation of a set of features of different degrees. In this repository, a fiber is represented as a dict where keys are degrees and values are the number of channels.
    • Multiplicity: The number of channels associated with a feature of a specific type.
    • Equivariance: A property where applying a symmetry transformation to the input and then the function produces the same result as applying the function and then the transformation.

  5. Install RoseTTAFold2NA

    main

    To install RoseTTAFold2NA, follow these steps to clone the repository, set up the Conda environment, install the SE(3)-Transformer dependency, and download the pre-trained weights.

    1. Clone the repository

    git clone https://github.com/uw-ipd/RoseTTAFold2NA.git
cd RoseTTAFold2NA

    2. Create Conda environment

    Use the provided RF2na-linux.yml to create the environment:

    conda env create -f RF2na-linux.yml

    3. Install SE(3)-Transformer

    You must install NVIDIA's SE(3)-Transformer using the version provided within this repository:

    conda activate RF2NA
cd SE3Transformer
pip install --no-cache-dir -r requirements.txt
python setup.py install
cd ..

    4. Download pre-trained weights

    Download the weights into the network directory:

    cd network
wget https://files.ipd.uw.edu/dimaio/RF2NA_apr23.tgz
tar xvfz RF2NA_apr23.tgz
ls weights/ # Verify the 1.1GB weights file is present
cd ..
    git clone https://github.com/uw-ipd/RoseTTAFold2NA.git
cd RoseTTAFold2NA
conda env create -f RF2na-linux.yml
conda activate RF2NA
cd SE3Transformer
pip install --no-cache-dir -r requirements.txt
python setup.py install
cd ..
cd network
wget https://files.ipd.uw.edu/dimaio/RF2NA_apr23.tgz
tar xvfz RF2NA_apr23.tgz
cd ..

  6. Quick Start Guide for SE(3)-Transformer

    main

    Follow these steps to clone the repository, build the Docker container, and run training and inference using default parameters on the QM9 dataset.

    1. Clone the repository
    2. Build the Docker container
    3. Start an interactive session (mounts a results directory for output)
    4. Run training via scripts/train.sh
    5. Run inference via scripts/predict.sh
    # 1. Clone
git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/DrugDiscovery/SE3Transformer

# 2. Build
docker build -t se3-transformer .

# 3. Run container
mkdir -p results
docker run -it --runtime=nvidia --shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 --rm -v ${PWD}/results:/results se3-transformer:latest

# 4. Train
bash scripts/train.sh

# 5. Predict
bash scripts/predict.sh

  7. Download sequence and structure databases

    main

    RoseTTAFold2NA requires several large databases for sequence and structure information. Follow these commands to download and prepare them:

    UniRef30 [46G]

    wget http://wwwuser.gwdg.de/~compbiol/uniclust/2020_06/UniRef30_2020_06_hhsuite.tar.gz
mkdir -p UniRef30_2020_06
tar xfz UniRef30_2020_06_hhsuite.tar.gz -C ./UniRef30_2020_06

    BFD [272G]

    wget https://bfd.mmseqs.com/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt.tar.gz
mkdir -p bfd
tar xfz bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt.tar.gz -C ./bfd

    Structure Templates

    wget https://files.ipd.uw.edu/pub/RoseTTAFold/pdb100_2021Mar03.tar.gz
tar xfz pdb100_2021Mar03.tar.gz

    RNA Databases

    Rfam [300M]

    mkdir -p RNA
cd RNA
wget ftp://ftp.ebi.ac.uk/pub/databases/Rfam/CURRENT/Rfam.full_region.gz
wget ftp://ftp.ebi.ac.uk/pub/databases/Rfam/CURRENT/Rfam.cm.gz
gunzip Rfam.cm.gz
cmpress Rfam.cm

    RNAcentral [12G]

    # Inside RNA directory
wget ftp://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/rfam/rfam_annotations.tsv.gz
wget ftp://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/id_mapping/id_mapping.tsv.gz
wget ftp://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/sequences/rnacentral_species_specific_ids.fasta.gz
../input_prep/reprocess_rnac.pl id_mapping.tsv.gz rfam_annotations.tsv.gz
gunzip -c rnacentral_species_specific_ids.fasta.gz | makeblastdb -in - -dbtype nucl -parse_seqids -out rnacentral.fasta -title "RNACentral"

    nt [151G]

    update_blastdb.pl --decompress nt
    # Example for UniRef30
wget http://wwwuser.gwdg.de/~compbiol/uniclust/2020_06/UniRef30_2020_06_hhsuite.tar.gz
mkdir -p UniRef30_2020_06
tar xfz UniRef30_2020_06_hhsuite.tar.gz -C ./UniRef30_2020_06

  8. Extend SE(3)-Transformer with custom datasets

    main

    To use the network on a new dataset, extend the DataModule class in se3_transformer/data_loading/data_module.py.

    Your custom collate function must return a tuple containing:

    • A (batched) DGLGraph object
    • A dictionary of node features ({'{degree}': tensor})
    • A dictionary of edge features ({'{degree}': tensor})
    • (Optional) Precomputed bases as a dictionary
    • Labels as a tensor

    After implementing the data module, you must modify training.py and inference.py to utilize your new class.

  10. How to enable Mixed Precision (AMP) training

    main

    This implementation supports Automatic Mixed Precision (AMP) using PyTorch's native package. AMP allows for FP16 training with FP32 master weights, providing significant speedups on NVIDIA Volta, Turing, and Ampere GPU architectures by leveraging Tensor Cores.

    To enable mixed precision during training or inference, use the --amp flag in your execution command.

    # Example usage (flag name provided in documentation)
python train_script.py --amp

  11. Troubleshoot memory allocation errors in Dataloader

    main

    If you encounter the error OSError: [Errno 12] Cannot allocate memory during Dataloader iterator creation (specifically during fork()), it is likely caused by the --precompute_bases flag.

    To resolve this without increasing RAM or Swap, disable bases precomputation by using one of the following methods:

    1. Remove the --precompute_bases flag from your command.
    2. Set the flag to false: --precompute_bases false.
    # Example fix: disable precompute_bases
python train.py --precompute_bases false ...

  12. Run training with Multi-GPU/Multi-node support

    main

    The training script supports the PyTorch elastic launcher. To train on all available GPUs with Automatic Mixed Precision (AMP) enabled, use torch.distributed.run to launch the se3_transformer.runtime.training module.

    python -m torch.distributed.run --nnodes=1 --nproc_per_node=gpu --module se3_transformer.runtime.training --amp