OneScience-Group/DiffDock

1#!/bin/bash2set -euo pipefail3export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/:$LD_LIBRARY_PATH"4export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"5export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH6 7SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)8EXAMPLE_DIR=$(cd "${SCRIPT_DIR}/.." && pwd)9REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd)10 11#source "${REPO_ROOT}/env.sh"12 13if [[ -n "${ROCM_PATH:-}" && -f "${ROCM_PATH}/cuda/env.sh" ]]; then14  source "${ROCM_PATH}/cuda/env.sh"15fi16 17export PYTHONPATH="${REPO_ROOT}/src:${REPO_ROOT}:${PYTHONPATH:-}"18export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}"19export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-${HIP_VISIBLE_DEVICES}}"20export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"21 22DIFFDOCK_DATA_ROOT="${DIFFDOCK_DATA_ROOT:-${ONESCIENCE_DATASETS_DIR}/diffdock}"23DIFFDOCK_DATASETS_DIR="${DIFFDOCK_DATASETS_DIR:-${DIFFDOCK_DATA_ROOT}/datasets}"24export TORCH_HOME="${TORCH_HOME:-${DIFFDOCK_DATA_ROOT}/torch_home}"25 26DATASET="${DATASET:-pdbbind}"27RUN_NAME="${RUN_NAME:-diffdock_${DATASET}_scnet}"28LOG_DIR="${LOG_DIR:-${EXAMPLE_DIR}/outputs/train}"29CACHE_PATH="${CACHE_PATH:-${EXAMPLE_DIR}/outputs/cache}"30CONFIG_PATH="${CONFIG_PATH:-${LOG_DIR}/${RUN_NAME}_train_config.yml}"31mkdir -p "${LOG_DIR}" "${CACHE_PATH}"32LOG_DIR=$(cd "${LOG_DIR}" && pwd)33CACHE_PATH=$(cd "${CACHE_PATH}" && pwd)34CONFIG_DIR=$(dirname "${CONFIG_PATH}")35CONFIG_NAME=$(basename "${CONFIG_PATH}")36mkdir -p "${CONFIG_DIR}"37CONFIG_PATH=$(cd "${CONFIG_DIR}" && pwd)/"${CONFIG_NAME}"38 39PDBBIND_DIR="${PDBBIND_DIR:-${DIFFDOCK_DATASETS_DIR}/PDBBind_processed}"40MOAD_DIR="${MOAD_DIR:-${DIFFDOCK_DATASETS_DIR}/BindingMOAD_2020_processed}"41SPLIT_TRAIN="${SPLIT_TRAIN:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_train}"42SPLIT_VAL="${SPLIT_VAL:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_val}"43 44DEVICE="${DEVICE:-auto}"45if [[ "$DEVICE" == "auto" ]]; then46    if python -c "import torch; print(torch.cuda.is_available())" | grep -q True; then47        DEVICE="cuda"48    else49        DEVICE="cpu"50    fi51fi52echo "****** DEVICE ******: $DEVICE"53 54SEED="${SEED:-0}"55BATCH_SIZE="${BATCH_SIZE:-4}"56N_EPOCHS="${N_EPOCHS:-10}"57LR="${LR:-0.001}"58NUM_WORKERS="${NUM_WORKERS:-1}"59NUM_DATALOADER_WORKERS="${NUM_DATALOADER_WORKERS:-0}"60LIMIT_COMPLEXES="${LIMIT_COMPLEXES:-null}"61SAVE_MODEL_FREQ="${SAVE_MODEL_FREQ:-null}"62VAL_INFERENCE_FREQ="${VAL_INFERENCE_FREQ:-null}"63TRAIN_INFERENCE_FREQ="${TRAIN_INFERENCE_FREQ:-null}"64WANDB="${WANDB:-false}"65NO_TORSION="${NO_TORSION:-false}"66NO_BATCH_NORM="${NO_BATCH_NORM:-true}"67 68PDBBIND_ESM_EMBEDDINGS_PATH="${PDBBIND_ESM_EMBEDDINGS_PATH:-}"69MOAD_ESM_EMBEDDINGS_PATH="${MOAD_ESM_EMBEDDINGS_PATH:-}"70MOAD_ESM_EMBEDDINGS_SEQUENCES_PATH="${MOAD_ESM_EMBEDDINGS_SEQUENCES_PATH:-}"71ESM_EMBEDDINGS_MODEL="${ESM_EMBEDDINGS_MODEL:-}"72 73yaml_value() {74  if [[ -z "${1:-}" || "${1}" == "null" ]]; then75    printf "null"76  else77    local value78    value=$(printf "%s" "$1" | sed "s/'/''/g")79    printf "'%s'" "$value"80  fi81}82 83yaml_bool() {84  if [[ "${1,,}" == "true" ]]; then85    printf "true"86  else87    printf "false"88  fi89}90 91cat > "${CONFIG_PATH}" <<EOF92runtime:93  run_name: $(yaml_value "${RUN_NAME}")94  log_dir: $(yaml_value "${LOG_DIR}")95  device: $(yaml_value "${DEVICE}")96  seed: ${SEED}97  cudnn_benchmark: false98  wandb: $(yaml_bool "${WANDB}")99  project: diffdock100 101data:102  dataset: $(yaml_value "${DATASET}")103  cache_path: $(yaml_value "${CACHE_PATH}")104  pdbbind_dir: $(yaml_value "${PDBBIND_DIR}")105  moad_dir: $(yaml_value "${MOAD_DIR}")106  split_train: $(yaml_value "${SPLIT_TRAIN}")107  split_val: $(yaml_value "${SPLIT_VAL}")108  protein_file: protein_processed109  limit_complexes: ${LIMIT_COMPLEXES}110  num_conformers: 1111  num_workers: ${NUM_WORKERS}112  num_dataloader_workers: ${NUM_DATALOADER_WORKERS}113  batch_size: ${BATCH_SIZE}114  pin_memory: false115  dataloader_drop_last: false116  remove_hs: true117  receptor_radius: 30118  c_alpha_max_neighbors: 10119  atom_radius: 5120  atom_max_neighbors: 8121  chain_cutoff: null122  max_lig_size: null123  matching_popsize: 20124  matching_maxiter: 20125  matching_tries: 1126  not_knn_only_graph: false127  include_miscellaneous_atoms: false128  all_atoms: false129  triple_training: false130  combined_training: false131  double_val: false132  train_multiplicity: 1133  val_multiplicity: 1134  max_receptor_size: null135  remove_promiscuous_targets: null136  min_ligand_size: 0137  unroll_clusters: false138  enforce_timesplit: false139  crop_beyond: 20140  moad_esm_embeddings_path: $(yaml_value "${MOAD_ESM_EMBEDDINGS_PATH}")141  pdbbind_esm_embeddings_path: $(yaml_value "${PDBBIND_ESM_EMBEDDINGS_PATH}")142  moad_esm_embeddings_sequences_path: $(yaml_value "${MOAD_ESM_EMBEDDINGS_SEQUENCES_PATH}")143  esm_embeddings_model: $(yaml_value "${ESM_EMBEDDINGS_MODEL}")144 145diffusion:146  no_torsion: $(yaml_bool "${NO_TORSION}")147  tr_sigma_min: 0.1148  tr_sigma_max: 30.0149  rot_sigma_min: 0.1150  rot_sigma_max: 1.65151  tor_sigma_min: 0.0314152  tor_sigma_max: 3.14153  sampling_alpha: 1.0154  sampling_beta: 1.0155  tr_weight: 0.33156  rot_weight: 0.33157  tor_weight: 0.33158  backbone_loss_weight: 0.0159  sidechain_loss_weight: 0.0160 161model:162  num_conv_layers: 2163  max_radius: 5.0164  scale_by_sigma: true165  norm_by_sigma: false166  ns: 16167  nv: 4168  distance_embed_dim: 32169  cross_distance_embed_dim: 32170  no_batch_norm: $(yaml_bool "${NO_BATCH_NORM}")171  use_second_order_repr: false172  cross_max_distance: 80173  dynamic_max_cross: false174  dropout: 0.0175  smooth_edges: false176  odd_parity: false177  embedding_type: sinusoidal178  sigma_embed_dim: 32179  embedding_scale: 1000180  no_aminoacid_identities: false181  sh_lmax: 2182  no_differentiate_convolutions: false183  tp_weights_layers: 2184  num_prot_emb_layers: 0185  reduce_pseudoscalars: false186  embed_also_ligand: true187  depthwise_convolution: false188  use_old_atom_encoder: false189 190optimization:191  n_epochs: ${N_EPOCHS}192  lr: ${LR}193  w_decay: 0.0194  scheduler: null195  scheduler_patience: 20196  lr_start_factor: 0.001197  warmup_dur: 4198  use_ema: true199  ema_rate: 0.999200  restart_dir: null201  restart_ckpt: last_model202  restart_lr: null203  pretrain_dir: null204  pretrain_ckpt: null205  save_model_freq: ${SAVE_MODEL_FREQ}206 207validation:208  test_sigma_intervals: false209  inference_samples: 1210  val_inference_freq: ${VAL_INFERENCE_FREQ}211  train_inference_freq: ${TRAIN_INFERENCE_FREQ}212  inference_steps: 20213  num_inference_complexes: 20214  inference_earlystop_metric: valinf_min_rmsds_lt2215  inference_secondary_metric: null216  inference_earlystop_goal: max217EOF218 219echo "DiffDock training config: ${CONFIG_PATH}"220echo "Dataset mode: ${DATASET}"221echo "PDBBind dir: ${PDBBIND_DIR}"222echo "MOAD dir: ${MOAD_DIR}"223echo "Train split: ${SPLIT_TRAIN}"224echo "Val split: ${SPLIT_VAL}"225echo "Cache path: ${CACHE_PATH}"226echo "Run output: ${LOG_DIR}/${RUN_NAME}"227echo "TORCH_HOME: ${TORCH_HOME}"228echo "Batch normalization disabled: ${NO_BATCH_NORM}"229 230cd "${REPO_ROOT}"231python "${SCRIPT_DIR}/train_diffdock.py" --config "${CONFIG_PATH}"232