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# echo "*** REPO_ROOT : ${REPO_ROOT}"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 26RUN_NAME="${RUN_NAME:-diffdock_pdbbind_smoke100_val20_cpu}"27LOG_DIR="${LOG_DIR:-${EXAMPLE_DIR}/outputs/train}"28CACHE_PATH="${CACHE_PATH:-${EXAMPLE_DIR}/outputs/cache_pdbbind_smoke100_val20}"29CONFIG_PATH="${CONFIG_PATH:-${LOG_DIR}/${RUN_NAME}_train_config.yml}"30 31DEVICE="${DEVICE:-cuda:0}"32BATCH_SIZE="${BATCH_SIZE:-2}"33N_EPOCHS="${N_EPOCHS:-1}"34NUM_WORKERS="${NUM_WORKERS:-1}"35NUM_DATALOADER_WORKERS="${NUM_DATALOADER_WORKERS:-0}"36LR="${LR:-0.001}"37SEED="${SEED:-0}"38 39PDBBIND_DIR="${PDBBIND_DIR:-${DIFFDOCK_DATASETS_DIR}/PDBBind_processed}"40SPLIT_TRAIN="${SPLIT_TRAIN:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_train_smoke100}"41SPLIT_VAL="${SPLIT_VAL:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_val_smoke20}"42 43mkdir -p "${LOG_DIR}" "${CACHE_PATH}"44LOG_DIR=$(cd "${LOG_DIR}" && pwd)45CACHE_PATH=$(cd "${CACHE_PATH}" && pwd)46CONFIG_DIR=$(dirname "${CONFIG_PATH}")47CONFIG_NAME=$(basename "${CONFIG_PATH}")48mkdir -p "${CONFIG_DIR}"49CONFIG_PATH=$(cd "${CONFIG_DIR}" && pwd)/"${CONFIG_NAME}"50 51yaml_value() {52  if [[ -z "${1:-}" || "${1}" == "null" ]]; then53    printf "null"54  else55    local value56    value=$(printf "%s" "$1" | sed "s/'/''/g")57    printf "'%s'" "$value"58  fi59}60 61cat > "${CONFIG_PATH}" <<EOF62runtime:63  run_name: $(yaml_value "${RUN_NAME}")64  log_dir: $(yaml_value "${LOG_DIR}")65  device: $(yaml_value "${DEVICE}")66  seed: ${SEED}67  cudnn_benchmark: false68  wandb: false69  project: diffdock70 71data:72  dataset: 'pdbbind'73  cache_path: $(yaml_value "${CACHE_PATH}")74  pdbbind_dir: $(yaml_value "${PDBBIND_DIR}")75  moad_dir: $(yaml_value "${PDBBIND_DIR}")76  split_train: $(yaml_value "${SPLIT_TRAIN}")77  split_val: $(yaml_value "${SPLIT_VAL}")78  protein_file: protein_processed79  limit_complexes: null80  num_conformers: 181  num_workers: ${NUM_WORKERS}82  num_dataloader_workers: ${NUM_DATALOADER_WORKERS}83  batch_size: ${BATCH_SIZE}84  pin_memory: false85  dataloader_drop_last: false86  remove_hs: true87  receptor_radius: 3088  c_alpha_max_neighbors: 1089  atom_radius: 590  atom_max_neighbors: 891  chain_cutoff: null92  max_lig_size: null93  matching_popsize: 2094  matching_maxiter: 2095  matching_tries: 196  not_knn_only_graph: false97  include_miscellaneous_atoms: false98  all_atoms: false99  triple_training: false100  combined_training: false101  double_val: false102  train_multiplicity: 1103  val_multiplicity: 1104  max_receptor_size: null105  remove_promiscuous_targets: null106  min_ligand_size: 0107  unroll_clusters: false108  enforce_timesplit: false109  crop_beyond: 20110  moad_esm_embeddings_path: null111  pdbbind_esm_embeddings_path: null112  moad_esm_embeddings_sequences_path: null113  esm_embeddings_model: null114 115diffusion:116  no_torsion: false117  tr_sigma_min: 0.1118  tr_sigma_max: 30.0119  rot_sigma_min: 0.1120  rot_sigma_max: 1.65121  tor_sigma_min: 0.0314122  tor_sigma_max: 3.14123  sampling_alpha: 1.0124  sampling_beta: 1.0125  tr_weight: 0.33126  rot_weight: 0.33127  tor_weight: 0.33128  backbone_loss_weight: 0.0129  sidechain_loss_weight: 0.0130 131model:132  num_conv_layers: 2133  max_radius: 5.0134  scale_by_sigma: true135  norm_by_sigma: false136  ns: 16137  nv: 4138  distance_embed_dim: 32139  cross_distance_embed_dim: 32140  no_batch_norm: true141  use_second_order_repr: false142  cross_max_distance: 80143  dynamic_max_cross: false144  dropout: 0.0145  smooth_edges: false146  odd_parity: false147  embedding_type: sinusoidal148  sigma_embed_dim: 32149  embedding_scale: 1000150  no_aminoacid_identities: false151  sh_lmax: 2152  no_differentiate_convolutions: false153  tp_weights_layers: 2154  num_prot_emb_layers: 0155  reduce_pseudoscalars: false156  embed_also_ligand: true157  depthwise_convolution: false158  use_old_atom_encoder: false159 160optimization:161  n_epochs: ${N_EPOCHS}162  lr: ${LR}163  w_decay: 0.0164  scheduler: null165  scheduler_patience: 20166  lr_start_factor: 0.001167  warmup_dur: 4168  use_ema: true169  ema_rate: 0.999170  restart_dir: null171  restart_ckpt: last_model172  restart_lr: null173  pretrain_dir: null174  pretrain_ckpt: null175  save_model_freq: null176 177validation:178  test_sigma_intervals: false179  inference_samples: 1180  val_inference_freq: null181  train_inference_freq: null182  inference_steps: 20183  num_inference_complexes: 20184  inference_earlystop_metric: valinf_min_rmsds_lt2185  inference_secondary_metric: null186  inference_earlystop_goal: max187EOF188 189echo "DiffDock smoke training config: ${CONFIG_PATH}"190echo "PDBBind dir: ${PDBBIND_DIR}"191echo "Train split: ${SPLIT_TRAIN}"192echo "Val split: ${SPLIT_VAL}"193echo "Cache path: ${CACHE_PATH}"194echo "Run output: ${LOG_DIR}/${RUN_NAME}"195echo "TORCH_HOME: ${TORCH_HOME}"196 197cd "${REPO_ROOT}"198python "${SCRIPT_DIR}/train_diffdock.py" --config "${CONFIG_PATH}"199 200