BranchesTags

Repository files navigation

nvDock - NVIDIA BioNeMo

Introduction

nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. Unlike the original DiffDock, which mainly represents proteins using Cα atoms, the all-atom version models the full atomic structure of the protein pocket. This allows the model to capture more detailed protein–ligand interactions and improves docking accuracy in complex binding environments.

The framework includes two components: a Score model and a Confidence model. The Score model generates multiple candidate binding poses for a ligand inside the target pocket through a reverse diffusion process, while the Confidence model ranks these poses to identify the most reliable binding configuration.

Score model is built on the neural network architecture introduced in the original DiffDock and is trained on curated protein–ligand complex datasets. By using all-atom protein representations together with pocket-focused training, the model can better learn local structural and chemical interactions within the binding site while maintaining efficient inference for docking tasks.

Confidence model is Heterogeneous SO(3)-Equivariant Graph Attention Network (HeteroEGA), a 3-dimensional SO(3)-equivariant heterogeneous graph neural network that has three stages:

  • (1) per-type node embedding with edge-degree initialization.
  • (2) interaction layers consisting of 3 HeteroEGA transformer blocks — each containing 9 Equivariant Graph Attention (EGA) modules over a heterogeneous graph (3 intra-type + 6 cross-type edges across ligand, receptor-residue, and receptor-atom nodes) and 3 per-type feed-forward networks with pre-norm residual connections
  • (3) an attention-weighted pooling readout over ligand nodes followed by an MLP confidence predictor.

Confidence model was developed based on: DiffDock-Pocket confidence architecture, replacing the Tensor-Product Convolutions with Equiformer v3 SO(2)-equivariant graph attention while preserving the original 9-way heterogeneous message-passing topology.

These models are available for both commercial and non-commercial use.

Environment Setup

Prerequisites

python ≥ 3.10.11
torch ≥ 2.6.0
nvidia-cublas-cu12 ≥ 12.5.0
cuequivariance ≥ 0.8.0
cuequivariance-torch ≥ 0.8.0

Data Preprocessing

This section describes how to perform the data preprocesing.

Score data

export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export HYDRA_FULL_ERROR=1
export PYTHONPATH=./src
python ${MOUNT_WORKING_PATH}/scripts/process_data.py \
name=${NAME} \
seed=${SEED} \
do_preprocessing=True \
do_training=False \
trainer.devices=1 \
trainer.num_nodes=1 \
trainer.max_epochs=1000 \
trainer.precision=32 \
data.num_workers=${NUM_WORKERS} \
data.cache_path=${CACHE_PATH} \
data.split_train=${SPLIT_TRAIN} \
data.split_val=${SPLIT_VAL} \
data.split_test=${SPLIT_TEST} \
data.process_embeddings=True \
data.all_atoms=True \
data.use_original_conformer_fallback=True \
protein_data.resume=True \
protein_data.protein_data_dir=${PROTEIN_DATA_DIR} \
protein_data.protein_ligand_csv=${PROTEIN_LIGAND_CSV} \
protein_data.output_fasta_file=${PDB_SEQUENCE_PATH} \
protein_data.esm_embeddings_path=${ESM_EMBEDINGS_PATH} \

Confidence data

export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=${MOUNT_WORKING_PATH}
NUM_WORKERS=32
MIN_NUM_WDS_TARS=100
DIR_HYDRA_CONFIG=${MOUNT_WORKING_PATH}/hydra_conf
PREFIX_CONFIG_SCORE_MODEL=score_model
SAMPLES_PER_COMPLEX=40
DENOISING_INFERENCE_STEPS=20
DOCKING_TYPE="pocket-docking"#Score model checkpoint
STATE_DICT_SCORE_INFER= # Data configuration
PROTEIN_LIGAND_CSV=
ESM_EMBEDDINGS_PATH=
DATA_DIR=
# Path to wds tar files of heterodata for score model infer
POCKET_HETERODATA_PATH=
DIR_INPUT_WDS_TARS_SCORE_INFER=${POCKET_HETERODATA_PATH}_train
# Split file (training/evaluation)
NAMES_SUBSET_SCORE_INFER=
# Path to full protein heterodata files
FULL_PROTEIN_HETERODATA=
# Output paths
DIR_OUTPUT_SCORE_INFER=./data_cache_wds_all_atoms/score_infer_outputs
DIR_OUTPUT_GRAPH=./data_cache_wds_all_atoms/complex_graph_allatom
DIR_OUTPUT_WDS_TARS=./data_cache_wds_all_atoms/training_set
python ${MOUNT_WORKING_PATH}/src/bionemo/diffdock/preprocessing/confidence_preprocess.py \
--protein_ligand_csv ${PROTEIN_LIGAND_CSV}\ 
--esm_embeddings_path ${ESM_EMBEDDINGS_PATH} \
--data_dir ${DATA_DIR} \
--dir_output_graph ${DIR_OUTPUT_GRAPH} \
--dir_output_wds_tars ${DIR_OUTPUT_WDS_TARS} \
--min_num_wds_tars ${MIN_NUM_WDS_TARS} \
--dir_hydra_config ${DIR_HYDRA_CONFIG} \
--prefix_config_score_model ${PREFIX_CONFIG_SCORE_MODEL} \
--state_dict_score_infer ${STATE_DICT_SCORE_INFER} \
--dir_input_wds_tars_score_infer ${DIR_INPUT_WDS_TARS_SCORE_INFER} \
--dir_output_score_infer ${DIR_OUTPUT_SCORE_INFER} \
--num_workers ${NUM_WORKERS} \
--names_subset_score_infer ${NAMES_SUBSET_SCORE_INFER} \
--samples_per_complex ${SAMPLES_PER_COMPLEX} \
--denoising_inference_steps ${DENOISING_INFERENCE_STEPS} \
--seed ${SEED} \
--docking_type ${DOCKING_TYPE}\ 
--full_protein_heterodata_path ${FULL_PROTEIN_HETERODATA}

Training

This section describes how to perform training from scratch.

train_score - Training score model

export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;

train_confidence - Training confidence model

export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_confidence.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_confidence_hetero \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
lightning_module.loss.rmsd_soft_label=True \
model.num_confidence_outputs=1 \
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;

Finetuning

DiffDock Pocket supports two fine-tuning methods, each with different trade-offs between computational requirements, parameter efficiency, and performance:

1. Full Fine-tuning: Updates all model parameters

Full fine-tuning updates all parameters in the model during training. This approach provides maximum flexibility since every weight can be adjusted to fit the target data.

Advantages

  • Generally provides the best performance
  • No constraints on which parts of the model can adapt
  • Can completely repurpose the model for significantly different data distributions

Limitations

  • Requires the most GPU memory
  • Can lead to overfitting, especially with small datasets
  • Needs longer training time
  • Requires storing the full model weights (~500MB+)

When to Use

  • When you have substantial training data (1000+ complexes)
  • When your computational resources are abundant
  • When your target domain differs significantly from the pre-training domain
  • When maximum performance is required

Example Command

export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=full
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;

2. Last Layer Fine-tuning: Updates only the final layers of the model

Last layer fine-tuning freezes most of the network but allows the final layers to be updated. This approach preserves the general features learned during pre-training while adapting the final prediction layers to the target task.

Advantages

  • Significantly less prone to overfitting than full fine-tuning
  • Requires less training data
  • Faster training than full fine-tuning
  • Preserves low-level features from pre-training

Limitations

  • Less flexibility than full fine-tuning
  • Cannot adapt the feature extraction components
  • May underperform if target domain differs greatly from pre-training

When to Use

  • When your dataset is moderate in size (100-1000 complexes)
  • When you want to fine-tune faster than full fine-tuning
  • When your target task is similar to the pre-training data
  • When you want to prevent overfitting on smaller datasets

Example Command

export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=last_layer
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;

Inference

This section describes how to perform inference using a trained model.

Supported inference modes include:

Single-sample inference – run docking inference for an individual protein–ligand pair

export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
complex_name=8DW5_FQ7 \
protein_path=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_protein.pdb \
ligand_description=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_ligand.sdf \
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight

Batch inference from CSV – perform inference on multiple samples specified in a CSV file

export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
protein_ligand_csv=${DATA_ROOT_DIR}/${SPLIT_CSV_NAME}.csv \ # Path to csv file
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight

Contribution Guidelines

  • This project is currently not accepting contributions.

About

nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. This all-atom version models the full atomic structure of the protein pocket, allows model to capture more detailed protein–ligand interactions, improves docking accuracy in complex binding environments.

Resources

Readme
License

Security policy

Security policy
Activity
Custom properties

Stars

2 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages

, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
BranchesTags
Repository files navigation
nvDock - NVIDIA BioNeMo
Introduction
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. Unlike the original DiffDock, which mainly represents proteins using Cα atoms, the all-atom version models the full atomic structure of the protein pocket. This allows the model to capture more detailed protein–ligand interactions and improves docking accuracy in complex binding environments.
The framework includes two components: a Score model and a Confidence model. The Score model generates multiple candidate binding poses for a ligand inside the target pocket through a reverse diffusion process, while the Confidence model ranks these poses to identify the most reliable binding configuration.
Score model is built on the neural network architecture introduced in the original DiffDock and is trained on curated protein–ligand complex datasets. By using all-atom protein representations together with pocket-focused training, the model can better learn local structural and chemical interactions within the binding site while maintaining efficient inference for docking tasks.
Confidence model is Heterogeneous SO(3)-Equivariant Graph Attention Network (HeteroEGA), a 3-dimensional SO(3)-equivariant heterogeneous graph neural network that has three stages:
  • (1) per-type node embedding with edge-degree initialization.
  • (2) interaction layers consisting of 3 HeteroEGA transformer blocks — each containing 9 Equivariant Graph Attention (EGA) modules over a heterogeneous graph (3 intra-type + 6 cross-type edges across ligand, receptor-residue, and receptor-atom nodes) and 3 per-type feed-forward networks with pre-norm residual connections
  • (3) an attention-weighted pooling readout over ligand nodes followed by an MLP confidence predictor.
Confidence model was developed based on: DiffDock-Pocket confidence architecture, replacing the Tensor-Product Convolutions with Equiformer v3 SO(2)-equivariant graph attention while preserving the original 9-way heterogeneous message-passing topology.
These models are available for both commercial and non-commercial use.
Environment Setup
Prerequisites
python ≥ 3.10.11
torch ≥ 2.6.0
nvidia-cublas-cu12 ≥ 12.5.0
cuequivariance ≥ 0.8.0
cuequivariance-torch ≥ 0.8.0
Data Preprocessing
This section describes how to perform the data preprocesing.
Score data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export HYDRA_FULL_ERROR=1
export PYTHONPATH=./src
python ${MOUNT_WORKING_PATH}/scripts/process_data.py \
name=${NAME} \
seed=${SEED} \
do_preprocessing=True \
do_training=False \
trainer.devices=1 \
trainer.num_nodes=1 \
trainer.max_epochs=1000 \
trainer.precision=32 \
data.num_workers=${NUM_WORKERS} \
data.cache_path=${CACHE_PATH} \
data.split_train=${SPLIT_TRAIN} \
data.split_val=${SPLIT_VAL} \
data.split_test=${SPLIT_TEST} \
data.process_embeddings=True \
data.all_atoms=True \
data.use_original_conformer_fallback=True \
protein_data.resume=True \
protein_data.protein_data_dir=${PROTEIN_DATA_DIR} \
protein_data.protein_ligand_csv=${PROTEIN_LIGAND_CSV} \
protein_data.output_fasta_file=${PDB_SEQUENCE_PATH} \
protein_data.esm_embeddings_path=${ESM_EMBEDINGS_PATH} \
Confidence data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=${MOUNT_WORKING_PATH}
NUM_WORKERS=32
MIN_NUM_WDS_TARS=100
DIR_HYDRA_CONFIG=${MOUNT_WORKING_PATH}/hydra_conf
PREFIX_CONFIG_SCORE_MODEL=score_model
SAMPLES_PER_COMPLEX=40
DENOISING_INFERENCE_STEPS=20
DOCKING_TYPE="pocket-docking"#Score model checkpoint
STATE_DICT_SCORE_INFER= # Data configuration
PROTEIN_LIGAND_CSV=
ESM_EMBEDDINGS_PATH=
DATA_DIR=
# Path to wds tar files of heterodata for score model infer
POCKET_HETERODATA_PATH=
DIR_INPUT_WDS_TARS_SCORE_INFER=${POCKET_HETERODATA_PATH}_train
# Split file (training/evaluation)
NAMES_SUBSET_SCORE_INFER=
# Path to full protein heterodata files
FULL_PROTEIN_HETERODATA=
# Output paths
DIR_OUTPUT_SCORE_INFER=./data_cache_wds_all_atoms/score_infer_outputs
DIR_OUTPUT_GRAPH=./data_cache_wds_all_atoms/complex_graph_allatom
DIR_OUTPUT_WDS_TARS=./data_cache_wds_all_atoms/training_set
python ${MOUNT_WORKING_PATH}/src/bionemo/diffdock/preprocessing/confidence_preprocess.py \
--protein_ligand_csv ${PROTEIN_LIGAND_CSV}\ 
--esm_embeddings_path ${ESM_EMBEDDINGS_PATH} \
--data_dir ${DATA_DIR} \
--dir_output_graph ${DIR_OUTPUT_GRAPH} \
--dir_output_wds_tars ${DIR_OUTPUT_WDS_TARS} \
--min_num_wds_tars ${MIN_NUM_WDS_TARS} \
--dir_hydra_config ${DIR_HYDRA_CONFIG} \
--prefix_config_score_model ${PREFIX_CONFIG_SCORE_MODEL} \
--state_dict_score_infer ${STATE_DICT_SCORE_INFER} \
--dir_input_wds_tars_score_infer ${DIR_INPUT_WDS_TARS_SCORE_INFER} \
--dir_output_score_infer ${DIR_OUTPUT_SCORE_INFER} \
--num_workers ${NUM_WORKERS} \
--names_subset_score_infer ${NAMES_SUBSET_SCORE_INFER} \
--samples_per_complex ${SAMPLES_PER_COMPLEX} \
--denoising_inference_steps ${DENOISING_INFERENCE_STEPS} \
--seed ${SEED} \
--docking_type ${DOCKING_TYPE}\ 
--full_protein_heterodata_path ${FULL_PROTEIN_HETERODATA}
Training
This section describes how to perform training from scratch.
train_score - Training score model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
train_confidence - Training confidence model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_confidence.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_confidence_hetero \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
lightning_module.loss.rmsd_soft_label=True \
model.num_confidence_outputs=1 \
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Finetuning
DiffDock Pocket supports two fine-tuning methods, each with different trade-offs between computational requirements, parameter efficiency, and performance:
1. Full Fine-tuning: Updates all model parameters
Full fine-tuning updates all parameters in the model during training. This approach provides maximum flexibility since every weight can be adjusted to fit the target data.
Advantages
  • Generally provides the best performance
  • No constraints on which parts of the model can adapt
  • Can completely repurpose the model for significantly different data distributions
Limitations
  • Requires the most GPU memory
  • Can lead to overfitting, especially with small datasets
  • Needs longer training time
  • Requires storing the full model weights (~500MB+)
When to Use
  • When you have substantial training data (1000+ complexes)
  • When your computational resources are abundant
  • When your target domain differs significantly from the pre-training domain
  • When maximum performance is required
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=full
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
2. Last Layer Fine-tuning: Updates only the final layers of the model
Last layer fine-tuning freezes most of the network but allows the final layers to be updated. This approach preserves the general features learned during pre-training while adapting the final prediction layers to the target task.
Advantages
  • Significantly less prone to overfitting than full fine-tuning
  • Requires less training data
  • Faster training than full fine-tuning
  • Preserves low-level features from pre-training
Limitations
  • Less flexibility than full fine-tuning
  • Cannot adapt the feature extraction components
  • May underperform if target domain differs greatly from pre-training
When to Use
  • When your dataset is moderate in size (100-1000 complexes)
  • When you want to fine-tune faster than full fine-tuning
  • When your target task is similar to the pre-training data
  • When you want to prevent overfitting on smaller datasets
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=last_layer
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Inference
This section describes how to perform inference using a trained model.
Supported inference modes include:
Single-sample inference – run docking inference for an individual protein–ligand pair
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
complex_name=8DW5_FQ7 \
protein_path=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_protein.pdb \
ligand_description=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_ligand.sdf \
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Batch inference from CSV – perform inference on multiple samples specified in a CSV file
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
protein_ligand_csv=${DATA_ROOT_DIR}/${SPLIT_CSV_NAME}.csv \ # Path to csv file
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Contribution Guidelines
  • This project is currently not accepting contributions.
About
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. This all-atom version models the full atomic structure of the protein pocket, allows model to capture more detailed protein–ligand interactions, improves docking accuracy in complex binding environments.
Resources
Readme
License
Security policy
Security policy
Activity
Custom properties
Stars
2 stars
Watchers
0 watching
Forks
0 forks
Report repository
Releases
Packages
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
BranchesTags
Repository files navigation
nvDock - NVIDIA BioNeMo
Introduction
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. Unlike the original DiffDock, which mainly represents proteins using Cα atoms, the all-atom version models the full atomic structure of the protein pocket. This allows the model to capture more detailed protein–ligand interactions and improves docking accuracy in complex binding environments.
The framework includes two components: a Score model and a Confidence model. The Score model generates multiple candidate binding poses for a ligand inside the target pocket through a reverse diffusion process, while the Confidence model ranks these poses to identify the most reliable binding configuration.
Score model is built on the neural network architecture introduced in the original DiffDock and is trained on curated protein–ligand complex datasets. By using all-atom protein representations together with pocket-focused training, the model can better learn local structural and chemical interactions within the binding site while maintaining efficient inference for docking tasks.
Confidence model is Heterogeneous SO(3)-Equivariant Graph Attention Network (HeteroEGA), a 3-dimensional SO(3)-equivariant heterogeneous graph neural network that has three stages:
  • (1) per-type node embedding with edge-degree initialization.
  • (2) interaction layers consisting of 3 HeteroEGA transformer blocks — each containing 9 Equivariant Graph Attention (EGA) modules over a heterogeneous graph (3 intra-type + 6 cross-type edges across ligand, receptor-residue, and receptor-atom nodes) and 3 per-type feed-forward networks with pre-norm residual connections
  • (3) an attention-weighted pooling readout over ligand nodes followed by an MLP confidence predictor.
Confidence model was developed based on: DiffDock-Pocket confidence architecture, replacing the Tensor-Product Convolutions with Equiformer v3 SO(2)-equivariant graph attention while preserving the original 9-way heterogeneous message-passing topology.
These models are available for both commercial and non-commercial use.
Environment Setup
Prerequisites
python ≥ 3.10.11
torch ≥ 2.6.0
nvidia-cublas-cu12 ≥ 12.5.0
cuequivariance ≥ 0.8.0
cuequivariance-torch ≥ 0.8.0
Data Preprocessing
This section describes how to perform the data preprocesing.
Score data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export HYDRA_FULL_ERROR=1
export PYTHONPATH=./src
python ${MOUNT_WORKING_PATH}/scripts/process_data.py \
name=${NAME} \
seed=${SEED} \
do_preprocessing=True \
do_training=False \
trainer.devices=1 \
trainer.num_nodes=1 \
trainer.max_epochs=1000 \
trainer.precision=32 \
data.num_workers=${NUM_WORKERS} \
data.cache_path=${CACHE_PATH} \
data.split_train=${SPLIT_TRAIN} \
data.split_val=${SPLIT_VAL} \
data.split_test=${SPLIT_TEST} \
data.process_embeddings=True \
data.all_atoms=True \
data.use_original_conformer_fallback=True \
protein_data.resume=True \
protein_data.protein_data_dir=${PROTEIN_DATA_DIR} \
protein_data.protein_ligand_csv=${PROTEIN_LIGAND_CSV} \
protein_data.output_fasta_file=${PDB_SEQUENCE_PATH} \
protein_data.esm_embeddings_path=${ESM_EMBEDINGS_PATH} \
Confidence data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=${MOUNT_WORKING_PATH}
NUM_WORKERS=32
MIN_NUM_WDS_TARS=100
DIR_HYDRA_CONFIG=${MOUNT_WORKING_PATH}/hydra_conf
PREFIX_CONFIG_SCORE_MODEL=score_model
SAMPLES_PER_COMPLEX=40
DENOISING_INFERENCE_STEPS=20
DOCKING_TYPE="pocket-docking"#Score model checkpoint
STATE_DICT_SCORE_INFER= # Data configuration
PROTEIN_LIGAND_CSV=
ESM_EMBEDDINGS_PATH=
DATA_DIR=
# Path to wds tar files of heterodata for score model infer
POCKET_HETERODATA_PATH=
DIR_INPUT_WDS_TARS_SCORE_INFER=${POCKET_HETERODATA_PATH}_train
# Split file (training/evaluation)
NAMES_SUBSET_SCORE_INFER=
# Path to full protein heterodata files
FULL_PROTEIN_HETERODATA=
# Output paths
DIR_OUTPUT_SCORE_INFER=./data_cache_wds_all_atoms/score_infer_outputs
DIR_OUTPUT_GRAPH=./data_cache_wds_all_atoms/complex_graph_allatom
DIR_OUTPUT_WDS_TARS=./data_cache_wds_all_atoms/training_set
python ${MOUNT_WORKING_PATH}/src/bionemo/diffdock/preprocessing/confidence_preprocess.py \
--protein_ligand_csv ${PROTEIN_LIGAND_CSV}\ 
--esm_embeddings_path ${ESM_EMBEDDINGS_PATH} \
--data_dir ${DATA_DIR} \
--dir_output_graph ${DIR_OUTPUT_GRAPH} \
--dir_output_wds_tars ${DIR_OUTPUT_WDS_TARS} \
--min_num_wds_tars ${MIN_NUM_WDS_TARS} \
--dir_hydra_config ${DIR_HYDRA_CONFIG} \
--prefix_config_score_model ${PREFIX_CONFIG_SCORE_MODEL} \
--state_dict_score_infer ${STATE_DICT_SCORE_INFER} \
--dir_input_wds_tars_score_infer ${DIR_INPUT_WDS_TARS_SCORE_INFER} \
--dir_output_score_infer ${DIR_OUTPUT_SCORE_INFER} \
--num_workers ${NUM_WORKERS} \
--names_subset_score_infer ${NAMES_SUBSET_SCORE_INFER} \
--samples_per_complex ${SAMPLES_PER_COMPLEX} \
--denoising_inference_steps ${DENOISING_INFERENCE_STEPS} \
--seed ${SEED} \
--docking_type ${DOCKING_TYPE}\ 
--full_protein_heterodata_path ${FULL_PROTEIN_HETERODATA}
Training
This section describes how to perform training from scratch.
train_score - Training score model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
train_confidence - Training confidence model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_confidence.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_confidence_hetero \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
lightning_module.loss.rmsd_soft_label=True \
model.num_confidence_outputs=1 \
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Finetuning
DiffDock Pocket supports two fine-tuning methods, each with different trade-offs between computational requirements, parameter efficiency, and performance:
1. Full Fine-tuning: Updates all model parameters
Full fine-tuning updates all parameters in the model during training. This approach provides maximum flexibility since every weight can be adjusted to fit the target data.
Advantages
  • Generally provides the best performance
  • No constraints on which parts of the model can adapt
  • Can completely repurpose the model for significantly different data distributions
Limitations
  • Requires the most GPU memory
  • Can lead to overfitting, especially with small datasets
  • Needs longer training time
  • Requires storing the full model weights (~500MB+)
When to Use
  • When you have substantial training data (1000+ complexes)
  • When your computational resources are abundant
  • When your target domain differs significantly from the pre-training domain
  • When maximum performance is required
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=full
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
2. Last Layer Fine-tuning: Updates only the final layers of the model
Last layer fine-tuning freezes most of the network but allows the final layers to be updated. This approach preserves the general features learned during pre-training while adapting the final prediction layers to the target task.
Advantages
  • Significantly less prone to overfitting than full fine-tuning
  • Requires less training data
  • Faster training than full fine-tuning
  • Preserves low-level features from pre-training
Limitations
  • Less flexibility than full fine-tuning
  • Cannot adapt the feature extraction components
  • May underperform if target domain differs greatly from pre-training
When to Use
  • When your dataset is moderate in size (100-1000 complexes)
  • When you want to fine-tune faster than full fine-tuning
  • When your target task is similar to the pre-training data
  • When you want to prevent overfitting on smaller datasets
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=last_layer
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Inference
This section describes how to perform inference using a trained model.
Supported inference modes include:
Single-sample inference – run docking inference for an individual protein–ligand pair
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
complex_name=8DW5_FQ7 \
protein_path=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_protein.pdb \
ligand_description=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_ligand.sdf \
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Batch inference from CSV – perform inference on multiple samples specified in a CSV file
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
protein_ligand_csv=${DATA_ROOT_DIR}/${SPLIT_CSV_NAME}.csv \ # Path to csv file
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Contribution Guidelines
  • This project is currently not accepting contributions.
About
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. This all-atom version models the full atomic structure of the protein pocket, allows model to capture more detailed protein–ligand interactions, improves docking accuracy in complex binding environments.
Resources
Readme
License
Security policy
Security policy
Activity
Custom properties
Stars
2 stars
Watchers
0 watching
Forks
0 forks
Report repository
Releases
Packages
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
BranchesTags
Repository files navigation
nvDock - NVIDIA BioNeMo
Introduction
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. Unlike the original DiffDock, which mainly represents proteins using Cα atoms, the all-atom version models the full atomic structure of the protein pocket. This allows the model to capture more detailed protein–ligand interactions and improves docking accuracy in complex binding environments.
The framework includes two components: a Score model and a Confidence model. The Score model generates multiple candidate binding poses for a ligand inside the target pocket through a reverse diffusion process, while the Confidence model ranks these poses to identify the most reliable binding configuration.
Score model is built on the neural network architecture introduced in the original DiffDock and is trained on curated protein–ligand complex datasets. By using all-atom protein representations together with pocket-focused training, the model can better learn local structural and chemical interactions within the binding site while maintaining efficient inference for docking tasks.
Confidence model is Heterogeneous SO(3)-Equivariant Graph Attention Network (HeteroEGA), a 3-dimensional SO(3)-equivariant heterogeneous graph neural network that has three stages:
  • (1) per-type node embedding with edge-degree initialization.
  • (2) interaction layers consisting of 3 HeteroEGA transformer blocks — each containing 9 Equivariant Graph Attention (EGA) modules over a heterogeneous graph (3 intra-type + 6 cross-type edges across ligand, receptor-residue, and receptor-atom nodes) and 3 per-type feed-forward networks with pre-norm residual connections
  • (3) an attention-weighted pooling readout over ligand nodes followed by an MLP confidence predictor.
Confidence model was developed based on: DiffDock-Pocket confidence architecture, replacing the Tensor-Product Convolutions with Equiformer v3 SO(2)-equivariant graph attention while preserving the original 9-way heterogeneous message-passing topology.
These models are available for both commercial and non-commercial use.
Environment Setup
Prerequisites
python ≥ 3.10.11
torch ≥ 2.6.0
nvidia-cublas-cu12 ≥ 12.5.0
cuequivariance ≥ 0.8.0
cuequivariance-torch ≥ 0.8.0
Data Preprocessing
This section describes how to perform the data preprocesing.
Score data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export HYDRA_FULL_ERROR=1
export PYTHONPATH=./src
python ${MOUNT_WORKING_PATH}/scripts/process_data.py \
name=${NAME} \
seed=${SEED} \
do_preprocessing=True \
do_training=False \
trainer.devices=1 \
trainer.num_nodes=1 \
trainer.max_epochs=1000 \
trainer.precision=32 \
data.num_workers=${NUM_WORKERS} \
data.cache_path=${CACHE_PATH} \
data.split_train=${SPLIT_TRAIN} \
data.split_val=${SPLIT_VAL} \
data.split_test=${SPLIT_TEST} \
data.process_embeddings=True \
data.all_atoms=True \
data.use_original_conformer_fallback=True \
protein_data.resume=True \
protein_data.protein_data_dir=${PROTEIN_DATA_DIR} \
protein_data.protein_ligand_csv=${PROTEIN_LIGAND_CSV} \
protein_data.output_fasta_file=${PDB_SEQUENCE_PATH} \
protein_data.esm_embeddings_path=${ESM_EMBEDINGS_PATH} \
Confidence data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=${MOUNT_WORKING_PATH}
NUM_WORKERS=32
MIN_NUM_WDS_TARS=100
DIR_HYDRA_CONFIG=${MOUNT_WORKING_PATH}/hydra_conf
PREFIX_CONFIG_SCORE_MODEL=score_model
SAMPLES_PER_COMPLEX=40
DENOISING_INFERENCE_STEPS=20
DOCKING_TYPE="pocket-docking"#Score model checkpoint
STATE_DICT_SCORE_INFER= # Data configuration
PROTEIN_LIGAND_CSV=
ESM_EMBEDDINGS_PATH=
DATA_DIR=
# Path to wds tar files of heterodata for score model infer
POCKET_HETERODATA_PATH=
DIR_INPUT_WDS_TARS_SCORE_INFER=${POCKET_HETERODATA_PATH}_train
# Split file (training/evaluation)
NAMES_SUBSET_SCORE_INFER=
# Path to full protein heterodata files
FULL_PROTEIN_HETERODATA=
# Output paths
DIR_OUTPUT_SCORE_INFER=./data_cache_wds_all_atoms/score_infer_outputs
DIR_OUTPUT_GRAPH=./data_cache_wds_all_atoms/complex_graph_allatom
DIR_OUTPUT_WDS_TARS=./data_cache_wds_all_atoms/training_set
python ${MOUNT_WORKING_PATH}/src/bionemo/diffdock/preprocessing/confidence_preprocess.py \
--protein_ligand_csv ${PROTEIN_LIGAND_CSV}\ 
--esm_embeddings_path ${ESM_EMBEDDINGS_PATH} \
--data_dir ${DATA_DIR} \
--dir_output_graph ${DIR_OUTPUT_GRAPH} \
--dir_output_wds_tars ${DIR_OUTPUT_WDS_TARS} \
--min_num_wds_tars ${MIN_NUM_WDS_TARS} \
--dir_hydra_config ${DIR_HYDRA_CONFIG} \
--prefix_config_score_model ${PREFIX_CONFIG_SCORE_MODEL} \
--state_dict_score_infer ${STATE_DICT_SCORE_INFER} \
--dir_input_wds_tars_score_infer ${DIR_INPUT_WDS_TARS_SCORE_INFER} \
--dir_output_score_infer ${DIR_OUTPUT_SCORE_INFER} \
--num_workers ${NUM_WORKERS} \
--names_subset_score_infer ${NAMES_SUBSET_SCORE_INFER} \
--samples_per_complex ${SAMPLES_PER_COMPLEX} \
--denoising_inference_steps ${DENOISING_INFERENCE_STEPS} \
--seed ${SEED} \
--docking_type ${DOCKING_TYPE}\ 
--full_protein_heterodata_path ${FULL_PROTEIN_HETERODATA}
Training
This section describes how to perform training from scratch.
train_score - Training score model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
train_confidence - Training confidence model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_confidence.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_confidence_hetero \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
lightning_module.loss.rmsd_soft_label=True \
model.num_confidence_outputs=1 \
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Finetuning
DiffDock Pocket supports two fine-tuning methods, each with different trade-offs between computational requirements, parameter efficiency, and performance:
1. Full Fine-tuning: Updates all model parameters
Full fine-tuning updates all parameters in the model during training. This approach provides maximum flexibility since every weight can be adjusted to fit the target data.
Advantages
  • Generally provides the best performance
  • No constraints on which parts of the model can adapt
  • Can completely repurpose the model for significantly different data distributions
Limitations
  • Requires the most GPU memory
  • Can lead to overfitting, especially with small datasets
  • Needs longer training time
  • Requires storing the full model weights (~500MB+)
When to Use
  • When you have substantial training data (1000+ complexes)
  • When your computational resources are abundant
  • When your target domain differs significantly from the pre-training domain
  • When maximum performance is required
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=full
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
2. Last Layer Fine-tuning: Updates only the final layers of the model
Last layer fine-tuning freezes most of the network but allows the final layers to be updated. This approach preserves the general features learned during pre-training while adapting the final prediction layers to the target task.
Advantages
  • Significantly less prone to overfitting than full fine-tuning
  • Requires less training data
  • Faster training than full fine-tuning
  • Preserves low-level features from pre-training
Limitations
  • Less flexibility than full fine-tuning
  • Cannot adapt the feature extraction components
  • May underperform if target domain differs greatly from pre-training
When to Use
  • When your dataset is moderate in size (100-1000 complexes)
  • When you want to fine-tune faster than full fine-tuning
  • When your target task is similar to the pre-training data
  • When you want to prevent overfitting on smaller datasets
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=last_layer
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Inference
This section describes how to perform inference using a trained model.
Supported inference modes include:
Single-sample inference – run docking inference for an individual protein–ligand pair
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
complex_name=8DW5_FQ7 \
protein_path=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_protein.pdb \
ligand_description=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_ligand.sdf \
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Batch inference from CSV – perform inference on multiple samples specified in a CSV file
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
protein_ligand_csv=${DATA_ROOT_DIR}/${SPLIT_CSV_NAME}.csv \ # Path to csv file
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Contribution Guidelines
  • This project is currently not accepting contributions.
About
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. This all-atom version models the full atomic structure of the protein pocket, allows model to capture more detailed protein–ligand interactions, improves docking accuracy in complex binding environments.
Resources
Readme
License
Security policy
Security policy
Activity
Custom properties
Stars
2 stars
Watchers
0 watching
Forks
0 forks
Report repository
Releases
Packages
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
BranchesTags
Repository files navigation
nvDock - NVIDIA BioNeMo
Introduction
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. Unlike the original DiffDock, which mainly represents proteins using Cα atoms, the all-atom version models the full atomic structure of the protein pocket. This allows the model to capture more detailed protein–ligand interactions and improves docking accuracy in complex binding environments.
The framework includes two components: a Score model and a Confidence model. The Score model generates multiple candidate binding poses for a ligand inside the target pocket through a reverse diffusion process, while the Confidence model ranks these poses to identify the most reliable binding configuration.
Score model is built on the neural network architecture introduced in the original DiffDock and is trained on curated protein–ligand complex datasets. By using all-atom protein representations together with pocket-focused training, the model can better learn local structural and chemical interactions within the binding site while maintaining efficient inference for docking tasks.
Confidence model is Heterogeneous SO(3)-Equivariant Graph Attention Network (HeteroEGA), a 3-dimensional SO(3)-equivariant heterogeneous graph neural network that has three stages:
  • (1) per-type node embedding with edge-degree initialization.
  • (2) interaction layers consisting of 3 HeteroEGA transformer blocks — each containing 9 Equivariant Graph Attention (EGA) modules over a heterogeneous graph (3 intra-type + 6 cross-type edges across ligand, receptor-residue, and receptor-atom nodes) and 3 per-type feed-forward networks with pre-norm residual connections
  • (3) an attention-weighted pooling readout over ligand nodes followed by an MLP confidence predictor.
Confidence model was developed based on: DiffDock-Pocket confidence architecture, replacing the Tensor-Product Convolutions with Equiformer v3 SO(2)-equivariant graph attention while preserving the original 9-way heterogeneous message-passing topology.
These models are available for both commercial and non-commercial use.
Environment Setup
Prerequisites
python ≥ 3.10.11
torch ≥ 2.6.0
nvidia-cublas-cu12 ≥ 12.5.0
cuequivariance ≥ 0.8.0
cuequivariance-torch ≥ 0.8.0
Data Preprocessing
This section describes how to perform the data preprocesing.
Score data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export HYDRA_FULL_ERROR=1
export PYTHONPATH=./src
python ${MOUNT_WORKING_PATH}/scripts/process_data.py \
name=${NAME} \
seed=${SEED} \
do_preprocessing=True \
do_training=False \
trainer.devices=1 \
trainer.num_nodes=1 \
trainer.max_epochs=1000 \
trainer.precision=32 \
data.num_workers=${NUM_WORKERS} \
data.cache_path=${CACHE_PATH} \
data.split_train=${SPLIT_TRAIN} \
data.split_val=${SPLIT_VAL} \
data.split_test=${SPLIT_TEST} \
data.process_embeddings=True \
data.all_atoms=True \
data.use_original_conformer_fallback=True \
protein_data.resume=True \
protein_data.protein_data_dir=${PROTEIN_DATA_DIR} \
protein_data.protein_ligand_csv=${PROTEIN_LIGAND_CSV} \
protein_data.output_fasta_file=${PDB_SEQUENCE_PATH} \
protein_data.esm_embeddings_path=${ESM_EMBEDINGS_PATH} \
Confidence data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=${MOUNT_WORKING_PATH}
NUM_WORKERS=32
MIN_NUM_WDS_TARS=100
DIR_HYDRA_CONFIG=${MOUNT_WORKING_PATH}/hydra_conf
PREFIX_CONFIG_SCORE_MODEL=score_model
SAMPLES_PER_COMPLEX=40
DENOISING_INFERENCE_STEPS=20
DOCKING_TYPE="pocket-docking"#Score model checkpoint
STATE_DICT_SCORE_INFER= # Data configuration
PROTEIN_LIGAND_CSV=
ESM_EMBEDDINGS_PATH=
DATA_DIR=
# Path to wds tar files of heterodata for score model infer
POCKET_HETERODATA_PATH=
DIR_INPUT_WDS_TARS_SCORE_INFER=${POCKET_HETERODATA_PATH}_train
# Split file (training/evaluation)
NAMES_SUBSET_SCORE_INFER=
# Path to full protein heterodata files
FULL_PROTEIN_HETERODATA=
# Output paths
DIR_OUTPUT_SCORE_INFER=./data_cache_wds_all_atoms/score_infer_outputs
DIR_OUTPUT_GRAPH=./data_cache_wds_all_atoms/complex_graph_allatom
DIR_OUTPUT_WDS_TARS=./data_cache_wds_all_atoms/training_set
python ${MOUNT_WORKING_PATH}/src/bionemo/diffdock/preprocessing/confidence_preprocess.py \
--protein_ligand_csv ${PROTEIN_LIGAND_CSV}\ 
--esm_embeddings_path ${ESM_EMBEDDINGS_PATH} \
--data_dir ${DATA_DIR} \
--dir_output_graph ${DIR_OUTPUT_GRAPH} \
--dir_output_wds_tars ${DIR_OUTPUT_WDS_TARS} \
--min_num_wds_tars ${MIN_NUM_WDS_TARS} \
--dir_hydra_config ${DIR_HYDRA_CONFIG} \
--prefix_config_score_model ${PREFIX_CONFIG_SCORE_MODEL} \
--state_dict_score_infer ${STATE_DICT_SCORE_INFER} \
--dir_input_wds_tars_score_infer ${DIR_INPUT_WDS_TARS_SCORE_INFER} \
--dir_output_score_infer ${DIR_OUTPUT_SCORE_INFER} \
--num_workers ${NUM_WORKERS} \
--names_subset_score_infer ${NAMES_SUBSET_SCORE_INFER} \
--samples_per_complex ${SAMPLES_PER_COMPLEX} \
--denoising_inference_steps ${DENOISING_INFERENCE_STEPS} \
--seed ${SEED} \
--docking_type ${DOCKING_TYPE}\ 
--full_protein_heterodata_path ${FULL_PROTEIN_HETERODATA}
Training
This section describes how to perform training from scratch.
train_score - Training score model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
train_confidence - Training confidence model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_confidence.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_confidence_hetero \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
lightning_module.loss.rmsd_soft_label=True \
model.num_confidence_outputs=1 \
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Finetuning
DiffDock Pocket supports two fine-tuning methods, each with different trade-offs between computational requirements, parameter efficiency, and performance:
1. Full Fine-tuning: Updates all model parameters
Full fine-tuning updates all parameters in the model during training. This approach provides maximum flexibility since every weight can be adjusted to fit the target data.
Advantages
  • Generally provides the best performance
  • No constraints on which parts of the model can adapt
  • Can completely repurpose the model for significantly different data distributions
Limitations
  • Requires the most GPU memory
  • Can lead to overfitting, especially with small datasets
  • Needs longer training time
  • Requires storing the full model weights (~500MB+)
When to Use
  • When you have substantial training data (1000+ complexes)
  • When your computational resources are abundant
  • When your target domain differs significantly from the pre-training domain
  • When maximum performance is required
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=full
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
2. Last Layer Fine-tuning: Updates only the final layers of the model
Last layer fine-tuning freezes most of the network but allows the final layers to be updated. This approach preserves the general features learned during pre-training while adapting the final prediction layers to the target task.
Advantages
  • Significantly less prone to overfitting than full fine-tuning
  • Requires less training data
  • Faster training than full fine-tuning
  • Preserves low-level features from pre-training
Limitations
  • Less flexibility than full fine-tuning
  • Cannot adapt the feature extraction components
  • May underperform if target domain differs greatly from pre-training
When to Use
  • When your dataset is moderate in size (100-1000 complexes)
  • When you want to fine-tune faster than full fine-tuning
  • When your target task is similar to the pre-training data
  • When you want to prevent overfitting on smaller datasets
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=last_layer
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Inference
This section describes how to perform inference using a trained model.
Supported inference modes include:
Single-sample inference – run docking inference for an individual protein–ligand pair
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
complex_name=8DW5_FQ7 \
protein_path=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_protein.pdb \
ligand_description=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_ligand.sdf \
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Batch inference from CSV – perform inference on multiple samples specified in a CSV file
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
protein_ligand_csv=${DATA_ROOT_DIR}/${SPLIT_CSV_NAME}.csv \ # Path to csv file
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Contribution Guidelines
  • This project is currently not accepting contributions.
About
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. This all-atom version models the full atomic structure of the protein pocket, allows model to capture more detailed protein–ligand interactions, improves docking accuracy in complex binding environments.
Resources
Readme
License
Security policy
Security policy
Activity
Custom properties
Stars
2 stars
Watchers
0 watching
Forks
0 forks
Report repository
Releases
Packages
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
BranchesTags
Repository files navigation
nvDock - NVIDIA BioNeMo
Introduction
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. Unlike the original DiffDock, which mainly represents proteins using Cα atoms, the all-atom version models the full atomic structure of the protein pocket. This allows the model to capture more detailed protein–ligand interactions and improves docking accuracy in complex binding environments.
The framework includes two components: a Score model and a Confidence model. The Score model generates multiple candidate binding poses for a ligand inside the target pocket through a reverse diffusion process, while the Confidence model ranks these poses to identify the most reliable binding configuration.
Score model is built on the neural network architecture introduced in the original DiffDock and is trained on curated protein–ligand complex datasets. By using all-atom protein representations together with pocket-focused training, the model can better learn local structural and chemical interactions within the binding site while maintaining efficient inference for docking tasks.
Confidence model is Heterogeneous SO(3)-Equivariant Graph Attention Network (HeteroEGA), a 3-dimensional SO(3)-equivariant heterogeneous graph neural network that has three stages:
  • (1) per-type node embedding with edge-degree initialization.
  • (2) interaction layers consisting of 3 HeteroEGA transformer blocks — each containing 9 Equivariant Graph Attention (EGA) modules over a heterogeneous graph (3 intra-type + 6 cross-type edges across ligand, receptor-residue, and receptor-atom nodes) and 3 per-type feed-forward networks with pre-norm residual connections
  • (3) an attention-weighted pooling readout over ligand nodes followed by an MLP confidence predictor.
Confidence model was developed based on: DiffDock-Pocket confidence architecture, replacing the Tensor-Product Convolutions with Equiformer v3 SO(2)-equivariant graph attention while preserving the original 9-way heterogeneous message-passing topology.
These models are available for both commercial and non-commercial use.
Environment Setup
Prerequisites
python ≥ 3.10.11
torch ≥ 2.6.0
nvidia-cublas-cu12 ≥ 12.5.0
cuequivariance ≥ 0.8.0
cuequivariance-torch ≥ 0.8.0
Data Preprocessing
This section describes how to perform the data preprocesing.
Score data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export HYDRA_FULL_ERROR=1
export PYTHONPATH=./src
python ${MOUNT_WORKING_PATH}/scripts/process_data.py \
name=${NAME} \
seed=${SEED} \
do_preprocessing=True \
do_training=False \
trainer.devices=1 \
trainer.num_nodes=1 \
trainer.max_epochs=1000 \
trainer.precision=32 \
data.num_workers=${NUM_WORKERS} \
data.cache_path=${CACHE_PATH} \
data.split_train=${SPLIT_TRAIN} \
data.split_val=${SPLIT_VAL} \
data.split_test=${SPLIT_TEST} \
data.process_embeddings=True \
data.all_atoms=True \
data.use_original_conformer_fallback=True \
protein_data.resume=True \
protein_data.protein_data_dir=${PROTEIN_DATA_DIR} \
protein_data.protein_ligand_csv=${PROTEIN_LIGAND_CSV} \
protein_data.output_fasta_file=${PDB_SEQUENCE_PATH} \
protein_data.esm_embeddings_path=${ESM_EMBEDINGS_PATH} \
Confidence data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=${MOUNT_WORKING_PATH}
NUM_WORKERS=32
MIN_NUM_WDS_TARS=100
DIR_HYDRA_CONFIG=${MOUNT_WORKING_PATH}/hydra_conf
PREFIX_CONFIG_SCORE_MODEL=score_model
SAMPLES_PER_COMPLEX=40
DENOISING_INFERENCE_STEPS=20
DOCKING_TYPE="pocket-docking"#Score model checkpoint
STATE_DICT_SCORE_INFER= # Data configuration
PROTEIN_LIGAND_CSV=
ESM_EMBEDDINGS_PATH=
DATA_DIR=
# Path to wds tar files of heterodata for score model infer
POCKET_HETERODATA_PATH=
DIR_INPUT_WDS_TARS_SCORE_INFER=${POCKET_HETERODATA_PATH}_train
# Split file (training/evaluation)
NAMES_SUBSET_SCORE_INFER=
# Path to full protein heterodata files
FULL_PROTEIN_HETERODATA=
# Output paths
DIR_OUTPUT_SCORE_INFER=./data_cache_wds_all_atoms/score_infer_outputs
DIR_OUTPUT_GRAPH=./data_cache_wds_all_atoms/complex_graph_allatom
DIR_OUTPUT_WDS_TARS=./data_cache_wds_all_atoms/training_set
python ${MOUNT_WORKING_PATH}/src/bionemo/diffdock/preprocessing/confidence_preprocess.py \
--protein_ligand_csv ${PROTEIN_LIGAND_CSV}\ 
--esm_embeddings_path ${ESM_EMBEDDINGS_PATH} \
--data_dir ${DATA_DIR} \
--dir_output_graph ${DIR_OUTPUT_GRAPH} \
--dir_output_wds_tars ${DIR_OUTPUT_WDS_TARS} \
--min_num_wds_tars ${MIN_NUM_WDS_TARS} \
--dir_hydra_config ${DIR_HYDRA_CONFIG} \
--prefix_config_score_model ${PREFIX_CONFIG_SCORE_MODEL} \
--state_dict_score_infer ${STATE_DICT_SCORE_INFER} \
--dir_input_wds_tars_score_infer ${DIR_INPUT_WDS_TARS_SCORE_INFER} \
--dir_output_score_infer ${DIR_OUTPUT_SCORE_INFER} \
--num_workers ${NUM_WORKERS} \
--names_subset_score_infer ${NAMES_SUBSET_SCORE_INFER} \
--samples_per_complex ${SAMPLES_PER_COMPLEX} \
--denoising_inference_steps ${DENOISING_INFERENCE_STEPS} \
--seed ${SEED} \
--docking_type ${DOCKING_TYPE}\ 
--full_protein_heterodata_path ${FULL_PROTEIN_HETERODATA}
Training
This section describes how to perform training from scratch.
train_score - Training score model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
train_confidence - Training confidence model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_confidence.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_confidence_hetero \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
lightning_module.loss.rmsd_soft_label=True \
model.num_confidence_outputs=1 \
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Finetuning
DiffDock Pocket supports two fine-tuning methods, each with different trade-offs between computational requirements, parameter efficiency, and performance:
1. Full Fine-tuning: Updates all model parameters
Full fine-tuning updates all parameters in the model during training. This approach provides maximum flexibility since every weight can be adjusted to fit the target data.
Advantages
  • Generally provides the best performance
  • No constraints on which parts of the model can adapt
  • Can completely repurpose the model for significantly different data distributions
Limitations
  • Requires the most GPU memory
  • Can lead to overfitting, especially with small datasets
  • Needs longer training time
  • Requires storing the full model weights (~500MB+)
When to Use
  • When you have substantial training data (1000+ complexes)
  • When your computational resources are abundant
  • When your target domain differs significantly from the pre-training domain
  • When maximum performance is required
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=full
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
2. Last Layer Fine-tuning: Updates only the final layers of the model
Last layer fine-tuning freezes most of the network but allows the final layers to be updated. This approach preserves the general features learned during pre-training while adapting the final prediction layers to the target task.
Advantages
  • Significantly less prone to overfitting than full fine-tuning
  • Requires less training data
  • Faster training than full fine-tuning
  • Preserves low-level features from pre-training
Limitations
  • Less flexibility than full fine-tuning
  • Cannot adapt the feature extraction components
  • May underperform if target domain differs greatly from pre-training
When to Use
  • When your dataset is moderate in size (100-1000 complexes)
  • When you want to fine-tune faster than full fine-tuning
  • When your target task is similar to the pre-training data
  • When you want to prevent overfitting on smaller datasets
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=last_layer
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Inference
This section describes how to perform inference using a trained model.
Supported inference modes include:
Single-sample inference – run docking inference for an individual protein–ligand pair
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
complex_name=8DW5_FQ7 \
protein_path=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_protein.pdb \
ligand_description=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_ligand.sdf \
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Batch inference from CSV – perform inference on multiple samples specified in a CSV file
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
protein_ligand_csv=${DATA_ROOT_DIR}/${SPLIT_CSV_NAME}.csv \ # Path to csv file
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Contribution Guidelines
  • This project is currently not accepting contributions.
About
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. This all-atom version models the full atomic structure of the protein pocket, allows model to capture more detailed protein–ligand interactions, improves docking accuracy in complex binding environments.
Resources
Readme
License
Security policy
Security policy
Activity
Custom properties
Stars
2 stars
Watchers
0 watching
Forks
0 forks
Report repository
Releases
Packages
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
BranchesTags
Repository files navigation
nvDock - NVIDIA BioNeMo
Introduction
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. Unlike the original DiffDock, which mainly represents proteins using Cα atoms, the all-atom version models the full atomic structure of the protein pocket. This allows the model to capture more detailed protein–ligand interactions and improves docking accuracy in complex binding environments.
The framework includes two components: a Score model and a Confidence model. The Score model generates multiple candidate binding poses for a ligand inside the target pocket through a reverse diffusion process, while the Confidence model ranks these poses to identify the most reliable binding configuration.
Score model is built on the neural network architecture introduced in the original DiffDock and is trained on curated protein–ligand complex datasets. By using all-atom protein representations together with pocket-focused training, the model can better learn local structural and chemical interactions within the binding site while maintaining efficient inference for docking tasks.
Confidence model is Heterogeneous SO(3)-Equivariant Graph Attention Network (HeteroEGA), a 3-dimensional SO(3)-equivariant heterogeneous graph neural network that has three stages:
  • (1) per-type node embedding with edge-degree initialization.
  • (2) interaction layers consisting of 3 HeteroEGA transformer blocks — each containing 9 Equivariant Graph Attention (EGA) modules over a heterogeneous graph (3 intra-type + 6 cross-type edges across ligand, receptor-residue, and receptor-atom nodes) and 3 per-type feed-forward networks with pre-norm residual connections
  • (3) an attention-weighted pooling readout over ligand nodes followed by an MLP confidence predictor.
Confidence model was developed based on: DiffDock-Pocket confidence architecture, replacing the Tensor-Product Convolutions with Equiformer v3 SO(2)-equivariant graph attention while preserving the original 9-way heterogeneous message-passing topology.
These models are available for both commercial and non-commercial use.
Environment Setup
Prerequisites
python ≥ 3.10.11
torch ≥ 2.6.0
nvidia-cublas-cu12 ≥ 12.5.0
cuequivariance ≥ 0.8.0
cuequivariance-torch ≥ 0.8.0
Data Preprocessing
This section describes how to perform the data preprocesing.
Score data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export HYDRA_FULL_ERROR=1
export PYTHONPATH=./src
python ${MOUNT_WORKING_PATH}/scripts/process_data.py \
name=${NAME} \
seed=${SEED} \
do_preprocessing=True \
do_training=False \
trainer.devices=1 \
trainer.num_nodes=1 \
trainer.max_epochs=1000 \
trainer.precision=32 \
data.num_workers=${NUM_WORKERS} \
data.cache_path=${CACHE_PATH} \
data.split_train=${SPLIT_TRAIN} \
data.split_val=${SPLIT_VAL} \
data.split_test=${SPLIT_TEST} \
data.process_embeddings=True \
data.all_atoms=True \
data.use_original_conformer_fallback=True \
protein_data.resume=True \
protein_data.protein_data_dir=${PROTEIN_DATA_DIR} \
protein_data.protein_ligand_csv=${PROTEIN_LIGAND_CSV} \
protein_data.output_fasta_file=${PDB_SEQUENCE_PATH} \
protein_data.esm_embeddings_path=${ESM_EMBEDINGS_PATH} \
Confidence data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=${MOUNT_WORKING_PATH}
NUM_WORKERS=32
MIN_NUM_WDS_TARS=100
DIR_HYDRA_CONFIG=${MOUNT_WORKING_PATH}/hydra_conf
PREFIX_CONFIG_SCORE_MODEL=score_model
SAMPLES_PER_COMPLEX=40
DENOISING_INFERENCE_STEPS=20
DOCKING_TYPE="pocket-docking"#Score model checkpoint
STATE_DICT_SCORE_INFER= # Data configuration
PROTEIN_LIGAND_CSV=
ESM_EMBEDDINGS_PATH=
DATA_DIR=
# Path to wds tar files of heterodata for score model infer
POCKET_HETERODATA_PATH=
DIR_INPUT_WDS_TARS_SCORE_INFER=${POCKET_HETERODATA_PATH}_train
# Split file (training/evaluation)
NAMES_SUBSET_SCORE_INFER=
# Path to full protein heterodata files
FULL_PROTEIN_HETERODATA=
# Output paths
DIR_OUTPUT_SCORE_INFER=./data_cache_wds_all_atoms/score_infer_outputs
DIR_OUTPUT_GRAPH=./data_cache_wds_all_atoms/complex_graph_allatom
DIR_OUTPUT_WDS_TARS=./data_cache_wds_all_atoms/training_set
python ${MOUNT_WORKING_PATH}/src/bionemo/diffdock/preprocessing/confidence_preprocess.py \
--protein_ligand_csv ${PROTEIN_LIGAND_CSV}\ 
--esm_embeddings_path ${ESM_EMBEDDINGS_PATH} \
--data_dir ${DATA_DIR} \
--dir_output_graph ${DIR_OUTPUT_GRAPH} \
--dir_output_wds_tars ${DIR_OUTPUT_WDS_TARS} \
--min_num_wds_tars ${MIN_NUM_WDS_TARS} \
--dir_hydra_config ${DIR_HYDRA_CONFIG} \
--prefix_config_score_model ${PREFIX_CONFIG_SCORE_MODEL} \
--state_dict_score_infer ${STATE_DICT_SCORE_INFER} \
--dir_input_wds_tars_score_infer ${DIR_INPUT_WDS_TARS_SCORE_INFER} \
--dir_output_score_infer ${DIR_OUTPUT_SCORE_INFER} \
--num_workers ${NUM_WORKERS} \
--names_subset_score_infer ${NAMES_SUBSET_SCORE_INFER} \
--samples_per_complex ${SAMPLES_PER_COMPLEX} \
--denoising_inference_steps ${DENOISING_INFERENCE_STEPS} \
--seed ${SEED} \
--docking_type ${DOCKING_TYPE}\ 
--full_protein_heterodata_path ${FULL_PROTEIN_HETERODATA}
Training
This section describes how to perform training from scratch.
train_score - Training score model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
train_confidence - Training confidence model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_confidence.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_confidence_hetero \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
lightning_module.loss.rmsd_soft_label=True \
model.num_confidence_outputs=1 \
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Finetuning
DiffDock Pocket supports two fine-tuning methods, each with different trade-offs between computational requirements, parameter efficiency, and performance:
1. Full Fine-tuning: Updates all model parameters
Full fine-tuning updates all parameters in the model during training. This approach provides maximum flexibility since every weight can be adjusted to fit the target data.
Advantages
  • Generally provides the best performance
  • No constraints on which parts of the model can adapt
  • Can completely repurpose the model for significantly different data distributions
Limitations
  • Requires the most GPU memory
  • Can lead to overfitting, especially with small datasets
  • Needs longer training time
  • Requires storing the full model weights (~500MB+)
When to Use
  • When you have substantial training data (1000+ complexes)
  • When your computational resources are abundant
  • When your target domain differs significantly from the pre-training domain
  • When maximum performance is required
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=full
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
2. Last Layer Fine-tuning: Updates only the final layers of the model
Last layer fine-tuning freezes most of the network but allows the final layers to be updated. This approach preserves the general features learned during pre-training while adapting the final prediction layers to the target task.
Advantages
  • Significantly less prone to overfitting than full fine-tuning
  • Requires less training data
  • Faster training than full fine-tuning
  • Preserves low-level features from pre-training
Limitations
  • Less flexibility than full fine-tuning
  • Cannot adapt the feature extraction components
  • May underperform if target domain differs greatly from pre-training
When to Use
  • When your dataset is moderate in size (100-1000 complexes)
  • When you want to fine-tune faster than full fine-tuning
  • When your target task is similar to the pre-training data
  • When you want to prevent overfitting on smaller datasets
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=last_layer
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Inference
This section describes how to perform inference using a trained model.
Supported inference modes include:
Single-sample inference – run docking inference for an individual protein–ligand pair
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
complex_name=8DW5_FQ7 \
protein_path=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_protein.pdb \
ligand_description=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_ligand.sdf \
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Batch inference from CSV – perform inference on multiple samples specified in a CSV file
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
protein_ligand_csv=${DATA_ROOT_DIR}/${SPLIT_CSV_NAME}.csv \ # Path to csv file
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Contribution Guidelines
  • This project is currently not accepting contributions.
About
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. This all-atom version models the full atomic structure of the protein pocket, allows model to capture more detailed protein–ligand interactions, improves docking accuracy in complex binding environments.
Resources
Readme
License
Security policy
Security policy
Activity
Custom properties
Stars
2 stars
Watchers
0 watching
Forks
0 forks
Report repository
Releases
Packages
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
BranchesTags
Repository files navigation
nvDock - NVIDIA BioNeMo
Introduction
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. Unlike the original DiffDock, which mainly represents proteins using Cα atoms, the all-atom version models the full atomic structure of the protein pocket. This allows the model to capture more detailed protein–ligand interactions and improves docking accuracy in complex binding environments.
The framework includes two components: a Score model and a Confidence model. The Score model generates multiple candidate binding poses for a ligand inside the target pocket through a reverse diffusion process, while the Confidence model ranks these poses to identify the most reliable binding configuration.
Score model is built on the neural network architecture introduced in the original DiffDock and is trained on curated protein–ligand complex datasets. By using all-atom protein representations together with pocket-focused training, the model can better learn local structural and chemical interactions within the binding site while maintaining efficient inference for docking tasks.
Confidence model is Heterogeneous SO(3)-Equivariant Graph Attention Network (HeteroEGA), a 3-dimensional SO(3)-equivariant heterogeneous graph neural network that has three stages:
  • (1) per-type node embedding with edge-degree initialization.
  • (2) interaction layers consisting of 3 HeteroEGA transformer blocks — each containing 9 Equivariant Graph Attention (EGA) modules over a heterogeneous graph (3 intra-type + 6 cross-type edges across ligand, receptor-residue, and receptor-atom nodes) and 3 per-type feed-forward networks with pre-norm residual connections
  • (3) an attention-weighted pooling readout over ligand nodes followed by an MLP confidence predictor.
Confidence model was developed based on: DiffDock-Pocket confidence architecture, replacing the Tensor-Product Convolutions with Equiformer v3 SO(2)-equivariant graph attention while preserving the original 9-way heterogeneous message-passing topology.
These models are available for both commercial and non-commercial use.
Environment Setup
Prerequisites
python ≥ 3.10.11
torch ≥ 2.6.0
nvidia-cublas-cu12 ≥ 12.5.0
cuequivariance ≥ 0.8.0
cuequivariance-torch ≥ 0.8.0
Data Preprocessing
This section describes how to perform the data preprocesing.
Score data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export HYDRA_FULL_ERROR=1
export PYTHONPATH=./src
python ${MOUNT_WORKING_PATH}/scripts/process_data.py \
name=${NAME} \
seed=${SEED} \
do_preprocessing=True \
do_training=False \
trainer.devices=1 \
trainer.num_nodes=1 \
trainer.max_epochs=1000 \
trainer.precision=32 \
data.num_workers=${NUM_WORKERS} \
data.cache_path=${CACHE_PATH} \
data.split_train=${SPLIT_TRAIN} \
data.split_val=${SPLIT_VAL} \
data.split_test=${SPLIT_TEST} \
data.process_embeddings=True \
data.all_atoms=True \
data.use_original_conformer_fallback=True \
protein_data.resume=True \
protein_data.protein_data_dir=${PROTEIN_DATA_DIR} \
protein_data.protein_ligand_csv=${PROTEIN_LIGAND_CSV} \
protein_data.output_fasta_file=${PDB_SEQUENCE_PATH} \
protein_data.esm_embeddings_path=${ESM_EMBEDINGS_PATH} \
Confidence data
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=${MOUNT_WORKING_PATH}
NUM_WORKERS=32
MIN_NUM_WDS_TARS=100
DIR_HYDRA_CONFIG=${MOUNT_WORKING_PATH}/hydra_conf
PREFIX_CONFIG_SCORE_MODEL=score_model
SAMPLES_PER_COMPLEX=40
DENOISING_INFERENCE_STEPS=20
DOCKING_TYPE="pocket-docking"#Score model checkpoint
STATE_DICT_SCORE_INFER= # Data configuration
PROTEIN_LIGAND_CSV=
ESM_EMBEDDINGS_PATH=
DATA_DIR=
# Path to wds tar files of heterodata for score model infer
POCKET_HETERODATA_PATH=
DIR_INPUT_WDS_TARS_SCORE_INFER=${POCKET_HETERODATA_PATH}_train
# Split file (training/evaluation)
NAMES_SUBSET_SCORE_INFER=
# Path to full protein heterodata files
FULL_PROTEIN_HETERODATA=
# Output paths
DIR_OUTPUT_SCORE_INFER=./data_cache_wds_all_atoms/score_infer_outputs
DIR_OUTPUT_GRAPH=./data_cache_wds_all_atoms/complex_graph_allatom
DIR_OUTPUT_WDS_TARS=./data_cache_wds_all_atoms/training_set
python ${MOUNT_WORKING_PATH}/src/bionemo/diffdock/preprocessing/confidence_preprocess.py \
--protein_ligand_csv ${PROTEIN_LIGAND_CSV}\ 
--esm_embeddings_path ${ESM_EMBEDDINGS_PATH} \
--data_dir ${DATA_DIR} \
--dir_output_graph ${DIR_OUTPUT_GRAPH} \
--dir_output_wds_tars ${DIR_OUTPUT_WDS_TARS} \
--min_num_wds_tars ${MIN_NUM_WDS_TARS} \
--dir_hydra_config ${DIR_HYDRA_CONFIG} \
--prefix_config_score_model ${PREFIX_CONFIG_SCORE_MODEL} \
--state_dict_score_infer ${STATE_DICT_SCORE_INFER} \
--dir_input_wds_tars_score_infer ${DIR_INPUT_WDS_TARS_SCORE_INFER} \
--dir_output_score_infer ${DIR_OUTPUT_SCORE_INFER} \
--num_workers ${NUM_WORKERS} \
--names_subset_score_infer ${NAMES_SUBSET_SCORE_INFER} \
--samples_per_complex ${SAMPLES_PER_COMPLEX} \
--denoising_inference_steps ${DENOISING_INFERENCE_STEPS} \
--seed ${SEED} \
--docking_type ${DOCKING_TYPE}\ 
--full_protein_heterodata_path ${FULL_PROTEIN_HETERODATA}
Training
This section describes how to perform training from scratch.
train_score - Training score model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
train_confidence - Training confidence model
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
python ${MOUNT_WORKING_PATH}/scripts/train_confidence.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_confidence_hetero \
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.train.data_path=[${PATH_TO_TRAIN_WDS}] \
data.val.data_path=${PATH_TO_VALIDATION_WDS} \
data.test.data_path=${PATH_TO_TEST_WDS} \
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
lightning_module.loss.rmsd_soft_label=True \
model.num_confidence_outputs=1 \
size_aware_batching.fraction_mem_cuda_cap=0.4 \
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Finetuning
DiffDock Pocket supports two fine-tuning methods, each with different trade-offs between computational requirements, parameter efficiency, and performance:
1. Full Fine-tuning: Updates all model parameters
Full fine-tuning updates all parameters in the model during training. This approach provides maximum flexibility since every weight can be adjusted to fit the target data.
Advantages
  • Generally provides the best performance
  • No constraints on which parts of the model can adapt
  • Can completely repurpose the model for significantly different data distributions
Limitations
  • Requires the most GPU memory
  • Can lead to overfitting, especially with small datasets
  • Needs longer training time
  • Requires storing the full model weights (~500MB+)
When to Use
  • When you have substantial training data (1000+ complexes)
  • When your computational resources are abundant
  • When your target domain differs significantly from the pre-training domain
  • When maximum performance is required
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=full
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
2. Last Layer Fine-tuning: Updates only the final layers of the model
Last layer fine-tuning freezes most of the network but allows the final layers to be updated. This approach preserves the general features learned during pre-training while adapting the final prediction layers to the target task.
Advantages
  • Significantly less prone to overfitting than full fine-tuning
  • Requires less training data
  • Faster training than full fine-tuning
  • Preserves low-level features from pre-training
Limitations
  • Less flexibility than full fine-tuning
  • Cannot adapt the feature extraction components
  • May underperform if target domain differs greatly from pre-training
When to Use
  • When your dataset is moderate in size (100-1000 complexes)
  • When you want to fine-tune faster than full fine-tuning
  • When your target task is similar to the pre-training data
  • When you want to prevent overfitting on smaller datasets
Example Command
export PYTORCH_CUDA_ALLOC_CONF=backend:native,expandable_segments:True;export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
FINETUNE_METHOD=last_layer
python ${MOUNT_WORKING_PATH}/scripts/finetune_score.py \
--config-path=${MOUNT_WORKING_PATH}/hydra_conf \
--config-name=train_score \
finetune_methods=${FINETUNE_METHOD} \
file_model_state_dict=${PRETRAINED_CHECKPOINT_PATH}\ # Path to pretrained checkpoint
name=${RUN_NAME}\ # experience name
seed=${SEED}\ # seed number
do_embedding_preprocessing=False \
do_preprocessing=False \
do_training=True \
trainer.devices=${GPUS_PER_NODE}\ # Number of GPUs per node
trainer.num_nodes=${NNODES}\ # Number of nodes
trainer.max_epochs=${NUM_EPOCHS}\ # Number of training epoch
trainer.precision=32 \ 
data.prefix_dir_wds=${PREFIX_CACHE_PATH}\ # Prefix path to webdataset
data.train.num_workers=${NUM_WORKERS}\ # Number of workers
data.train.micro_batch_size=${MICRO_BATCH_SIZE}\ # Number of batchsize per gpu
data.train.num_samples=${NUM_TRAIN_SAMPLES}\ # Number of training samples
data.val.num_samples=${NUM_VAL_SAMPLES}\ # Number of validation samples
data.test.num_samples=${NUM_TEST_SAMPLES}\ # Number of testing samples
model.all_atoms=True \ # Enable all-atom Model
exp_manager.exp_dir=${EXPERIMENT_PATH}\ # Path to save experiments
exp_manager.create_checkpoint_callback=True ;
Inference
This section describes how to perform inference using a trained model.
Supported inference modes include:
Single-sample inference – run docking inference for an individual protein–ligand pair
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
complex_name=8DW5_FQ7 \
protein_path=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_protein.pdb \
ligand_description=data/posebusters_benchmark_set/8DW5_FQ7/8DW5_FQ7_ligand.sdf \
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Batch inference from CSV – perform inference on multiple samples specified in a CSV file
export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
export BIONEMO_DIFFDOCK_PATH=MOUNT_WORKING_PATH # Path to source code
SAVE_DIR=#path to save predictions
mkdir -p ${SAVE_DIR}; python ${MOUNT_WORKING_PATH}/scripts/infer.py \
seed=${SEED} \
out_dir=${SAVE_DIR} \
protein_ligand_csv=${DATA_ROOT_DIR}/${SPLIT_CSV_NAME}.csv \ # Path to csv file
docking_type=pocket-docking \ # Set docking type to Pocket docking
++score_infer.file_model_state_dict=${CKPT_CHECKPOINT_PATH}\ # Path to score model weight
++score.model.all_atoms=True \ # Enable all-atom Model
++score_infer.data.all_atoms=True \ # Process data as all-atom
++score_infer.samples_per_complex=${SAMPLES_PER_COMPLEX}\ # Number of generated poses per complex
++confidence.model.num_confidence_outputs=1 \
++confidence_infer.backbone_type=equiformer_v3_hetero \ # Use EGA confidence model
++equiformer_v3_hetero_confidence.model.num_confidence_outputs=1 \
++confidence_infer.file_model_state_dict=${CONFIDENCE_PATH};# Path to confidence model weight
Contribution Guidelines
  • This project is currently not accepting contributions.
About
nvDock is a diffusion-based model for molecular pocket docking, designed for cases where the protein binding pocket is known in advance. This all-atom version models the full atomic structure of the protein pocket, allows model to capture more detailed protein–ligand interactions, improves docking accuracy in complex binding environments.
Resources
Readme
License
Security policy
Security policy
Activity
Custom properties
Stars
2 stars
Watchers
0 watching
Forks
0 forks
Report repository
Releases
Packages
Contributors
Languages