File size: 2,241 Bytes
9d6a2a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | #!/usr/bin/env bash
set -euo pipefail
source ${ROCM_PATH}/cuda/env.sh
export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
echo ${ROOT_DIR}
source ${ROOT_DIR}/env.sh
cd "${ROOT_DIR}"
export PYTHONPATH="${ROOT_DIR}/src:${PYTHONPATH:-}"
if [[ -n "${CONDA_PREFIX:-}" ]]; then
export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib:${CONDA_PREFIX}/lib/python3.11/site-packages/torch/lib:${CONDA_PREFIX}/lib/python3.11/site-packages/fastpt/torch/lib:/opt/dtk/.hyhal/rocm_smi/lib:${LD_LIBRARY_PATH:-}"
fi
GENSCORE_DATA_DIR="${GENSCORE_DATA_DIR:-${ONESCIENCE_DATASETS_DIR}/GenScore/genscore_data/rtmscore_s}"
GENSCORE_DATA_PREFIX="${GENSCORE_DATA_PREFIX:-v2020_train}"
GENSCORE_ENCODER="${GENSCORE_ENCODER:-gatedgcn}"
GENSCORE_MODEL_PATH="${GENSCORE_MODEL_PATH:-${SCRIPT_DIR}/genscore_${GENSCORE_ENCODER}_full_3000.pth}"
GENSCORE_NUM_EPOCHS="${GENSCORE_NUM_EPOCHS:-3000}"
GENSCORE_BATCH_SIZE="${GENSCORE_BATCH_SIZE:-64}"
GENSCORE_NUM_WORKERS="${GENSCORE_NUM_WORKERS:-8}"
GENSCORE_VALNUM="${GENSCORE_VALNUM:-1500}"
GENSCORE_PATIENCE="${GENSCORE_PATIENCE:-70}"
for suffix in ids.npy lig.pt prot.pt; do
path="${GENSCORE_DATA_DIR}/${GENSCORE_DATA_PREFIX}_${suffix}"
if [[ ! -f "${path}" ]]; then
echo "Missing required training file: ${path}" >&2
exit 1
fi
done
mkdir -p "$(dirname "${GENSCORE_MODEL_PATH}")"
echo "GenScore full training"
echo " data_dir: ${GENSCORE_DATA_DIR}"
echo " data_prefix: ${GENSCORE_DATA_PREFIX}"
echo " model_path: ${GENSCORE_MODEL_PATH}"
echo " encoder: ${GENSCORE_ENCODER}"
echo " num_epochs: ${GENSCORE_NUM_EPOCHS}"
echo " batch_size: ${GENSCORE_BATCH_SIZE}"
echo " num_workers: ${GENSCORE_NUM_WORKERS}"
echo " valnum: ${GENSCORE_VALNUM}"
echo " patience: ${GENSCORE_PATIENCE}"
python ${SCRIPT_DIR}/train_genscore.py \
--data_dir "${GENSCORE_DATA_DIR}" \
--data_prefix "${GENSCORE_DATA_PREFIX}" \
--model_path "${GENSCORE_MODEL_PATH}" \
--num_epochs "${GENSCORE_NUM_EPOCHS}" \
--batch_size "${GENSCORE_BATCH_SIZE}" \
--num_workers "${GENSCORE_NUM_WORKERS}" \
--valnum "${GENSCORE_VALNUM}" \
--patience "${GENSCORE_PATIENCE}" \
--encoder "${GENSCORE_ENCODER}"
|