Zenn
🍊

DiffDock実行環境の構築 (ver. 2023-02-24) | Apptainer (on WSL2 & with GPU)

2023/03/18に公開・約7,700字
WSL 2
#
container
#
bioinformatics
singularity
#
apptainer
tech

※ コードブロック における 🍊 は、プロンプトを表している

🍊 はじめに

以前に、DiffDock というdiffusion modelに基づく分子ドッキングツールの環境構築について記事を書いた (記事は削除済み)。

開発者であるCorso, G.氏のツイートによると、ICLR (The International Conference on Learning Representations) に論文がアクセプトされたそうで、それに伴い arXiv のプレプリントの更新と新しいコードを公開してくれたらしい (感謝！)。
arXiv | DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking
GitHub | DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking

コードの変更箇所について細かいところまでは追えていないが、大きな変更としては ESM embedding generation と ESMFold structure prediction の部分を DiffDock のコードに内包化してくれたようだ (結果として実行時の手間が減った)。

本記事では、2023-02-24版の DiffDock 実行環境構築 (with Apptainer ) について紹介します ( Apptainer 使わずに環境構築したい人の参考にもなるかと思います)。

以下、本記事を執筆した際の環境。

specifications
CPU Ryzen 7 5700X (AMD)
グラフィックボード GeForce RTX 2080 Ti VENTUS GP OC (MSI)
メモリ DDR4-3200 16 GB×2 (Crucial)
version
Windows 10 22H2
NVIDIA グラフィクス ドライバー 528.49
WSL2 1.0.3.0
AlmaLinux 8 on WLS2 8.7
Apptainer 1.1.5-2.el8

なお、WSL2 上の AlmaLinux 8 に Apptainer 環境を構築する手順については、 以下の記事をご参照ください。
Zenn | AlmaLinux 8 on WSL2にApptainer環境 (with GPU) を構築する

🍊 DiffDock実行コンテナの構築

基本的には DiffDock 公式のSetup Environmentに従う。
GitHub | DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking

ただし、公式では conda install を用いて PyTorch のインストールをしているが、私の環境ではうまくいかなかったため pip install を代わりに用いることにした ( PyTorch 自体は conda でインストールできたが、他のパッケージを pip でインストールする際に問題が生じた)。

以上を踏まえて、任意のディレクトリ (例として ~/apps/apptainer_image/def_file ) にコンテナイメージビルドのための DiffDock.def を作成する。

~/apps/apptainer_image/def_file/DiffDock.def
BootStrap: docker
From: nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04

%post
  # localtime関係のWARNINGに対処
  touch /etc/localtime
  
  # リポジトリを富山大学に変更
  sed -i.bak -e 's|archive.ubuntu.com|ubuntutym.u-toyama.ac.jp|g' /etc/apt/sources.list

  # DiffDockに必要なパッケージのインストール (冗長)
  apt update && apt upgrade -y
  apt install -y build-essential git curl
  
  # localeにen_US.UTF-8を追加する (WSL2だけで使用するなら不要)
  apt install -y locales
  locale-gen en_US.UTF-8
  
  # Pyenvを/usr/local/appsにインストール
  git clone https://github.com/yyuu/pyenv.git /usr/local/apps/pyenv
  export PYENV_ROOT="/usr/local/apps/pyenv"
  export PATH="${PYENV_ROOT}/bin:${PATH}"

  # PyenvでMiniforgeをインストールし、MiniforgeにPATHを通す
  pyenv install --list
  pyenv install miniforge3-22.9.0-3
  pyenv global miniforge3-22.9.0-3
  pyenv versions
  export MINIFORGE3_ROOT="${PYENV_ROOT}/versions/miniforge3-22.9.0-3"
  export PATH="${MINIFORGE3_ROOT}/bin:${PATH}"

  # conda自身のアップデート
  conda update -n base -c conda-forge --override-channels conda -y

  # DiffDock用のconda環境 (diffdock-conda) を作成しPATHを通す (activateのかわり)
  conda create -n diffdock-conda -c conda-forge --override-channels python=3.8 -y
  export PATH="${MINIFORGE3_ROOT}/envs/diffdock-conda/bin:${PATH}"

  # DiffDockに必要なパッケージをconda環境diffdock-condaにインストールする
  python3 -m pip install --upgrade pip
  python3 -m pip install torch torchvision torchaudio

  python3 -m pip install \
    torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric \
    -f https://data.pyg.org/whl/torch-1.13.0+cu117.html
  
  python3 -m pip install \
    PyYAML scipy "networkx[default]" biopython rdkit-pypi \
    e3nn spyrmsd pandas biopandas
  
  python3 -m pip install "fair-esm[esmfold]"
  python3 -m pip install \
    'dllogger @ git+https://github.com/NVIDIA/dllogger.git'
  python3 -m pip install \
    'openfold @ git+https://github.com/aqlaboratory/openfold.git@4b41059694619831a7db195b7e0988fc4ff3a307'
  
  # conda環境diffdock-condaにインストールしたパッケージの由来 (channel) を確認する
  conda list -n diffdock-conda

  # GitHubから/usr/local/appsにDiffDockを複製する
  git clone https://github.com/gcorso/DiffDock.git /usr/local/apps/DiffDock
  
  # ESM2で用いるpretrained modelを
  # コンテナ内 (/usr/local/apps/torch/hub/checkpoints) にあらかじめダウンロードしておく
  mkdir -p /usr/local/apps/torch/hub/checkpoints
  cd /usr/local/apps/torch/hub/checkpoints
  curl -O -L https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt
  curl -O -L https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t33_650M_UR50D-contact-regression.pt

%environment
  # DiffDock実行のための環境設定
  export PYENV_ROOT="/usr/local/apps/pyenv"
  export MINIFORGE3_ROOT="${PYENV_ROOT}/versions/miniforge3-22.9.0-3"
  export PATH="${MINIFORGE3_ROOT}/envs/diffdock-conda/bin:${PATH}"
  export DIFFDOCK_HOME="/usr/local/apps/DiffDock"
  export PYTHONPATH="${DIFFDOCK_HOME}"
  
  # TORCH_HOMEを定義しておくことで
  # コンテナ内にダウンロードした.ptファイルをESM2が使うようにする
  export TORCH_HOME="/usr/local/apps/torch"

%runscript
  # コンテナ実行時のスクリプトを定義
  # "$@"は追加の引数を利用するための設定
  python3 -m inference \
    --model_dir ${DIFFDOCK_HOME}/workdir/paper_score_model \
    --confidence_model_dir ${DIFFDOCK_HOME}/workdir/paper_confidence_model \
    "$@"

注意点は以下の2つ

  • DiffDock 実行時に ESM2 が使用するpretrained modelをコンテナビルドの段階でコンテナ内にダウンロードし、ダウンロード先を ${TORCH_HOME} として定義しておく ( %environment に記載 )
    これにより DiffDock 初回実行時、~/.cache/torch/hub/checkpoints (コンテナ外) に .pt ファイル がダウンロードされないようにしている
  • %runscript に --model_dir および --confidence_model_dir は記載しておくと良い ( DiffDock でドッキングする際、基本的に毎回指定する必要があるため)

作成した DiffDock.def を基に、DiffDock 実行コンテナ DiffDock.sif を以下のコマンドでビルドする。

🍊 cd ~/apps/apptainer_image
🍊 apptainer build ./DiffDock.sif ./def_file/DiffDock.def

🍊 DiffDockの動作テスト

DiffDock を用いたSARS-CoV-2 3CL Protease

蛋白質インプットファイル：
~/Documents/diffdock/7vu6/input/7vu6_A_noHET.pdb
※ PDB code: 7VU6 (with chain A without HETATM)

リガンドインプットファイル：
~/Documents/diffdock/7vu6/input/ensitrelvir_2D_162533924.sdf
※ PubChem CID: 162533924

DiffDock 実行スクリプト：
~/Documents/diffdock/7vu6/diffdock_run.sh

~/Documents/diffdock/7vu6/diffdock_run.sh
#!/usr/bin/bash

# diffdock_run.shを置いた場所を以降の起点とする
DOCKING_DIR="$(cd $(dirname $0) && pwd)"
cd ${DOCKING_DIR}

# 蛋白質およびリガンドのインプットファイルの場所 (都度修正)
protein_path="${DOCKING_DIR}/input/7vu6_A_noHET.pdb"
ligand_path="${DOCKING_DIR}/input/ensitrelvir_2D_162533924.sdf"

# 予測結果が出力されるディレクトリ名 (任意)
complex_name="7vu6_ensitrelvir"

# DiffDock.sifの格納場所 (適宜修正)
APPTAINER_IMAGE="${HOME}/apps/apptainer_image"

# DiffDockの実行 (各オプションに指定した数値は公式例)
apptainer run --nvccli ${APPTAINER_IMAGE}/DiffDock.sif \
  --protein_path ${protein_path} \
  --ligand ${ligand_path} \
  --complex_name ${complex_name} \
  --out_dir ${DOCKING_DIR}/results \
  --inference_steps 20 \
  --samples_per_complex 40 \
  --batch_size 10 \
  --actual_steps 18 \
  --no_final_step_noise

# 結果として不要 (そうな) ファイルの削除
rm -f ${DOCKING_DIR}/.*.npy

ESM2 による処理が DiffDock に内包化されたため、前回記事と比較して実行スクリプトがシンプルになった。

外部ディスクのファイルあるいは WSL2 からWindows側のファイルについて計算したい場合などは、apptainer run のオプション --bind を利用する。
例えば WSL2 においてWindowsのDドライブは /mnt/d にマウントされているので、apptainer run --bind /mnt/d:/mnt/d とすればよい。

作成した diffdock_run.sh を以下のように実行する。

# DiffDockの実行
🍊 cd ~/Documents/diffdock/7vu6
🍊 chmod u+x ./diffdock_run.sh
🍊 ./diffdock_run

👇 DiffDock を用いたドッキングシミュレーション結果
Green🍵: Crystal structure (PDB code: 7VU6)
Orange🍊: DiffDock prediction (top score)

Discussion

ログインするとコメントできます