AlphaFold2の最初の公開PyTorchバージョンが複製可能になりました。コロンビア大学のオープンソースで、1,000以上のスターが付いています。

ちょうど今、コロンビア大学のシステム生物学助教授であるモハメッド・アルクライシ氏が、AlphaFold2 のトレーニング可能な PyTorch 再現バージョンである OpenFold と呼ばれるモデルをゼロからトレーニングしたことを Twitter で発表しました。モハメッド・アルクライシ氏は、これが一般に公開される最初のAlphaFold2の複製品であるとも述べた。

AlphaFold2 は、原子レベルの精度でタンパク質構造を定期的に予測できます。技術的には、多重配列アライメントとディープラーニングアルゴリズム設計を活用し、タンパク質構造に関する物理的および生物学的知識を組み合わせて、予測結果を向上させます。タンパク質構造予測の2/3という優れた成果を達成し、昨年Nature誌に掲載されました。さらに驚くべきことは、DeepMind チームがモデルをオープンソース化しただけでなく、AlphaFold2 の予測データを無料のオープンデータセットにしたことです。

ただし、オープンソースだからといって、それが使用可能である、または使いやすいというわけではありません。実際、AlphaFold2 ソフトウェアシステムの展開は非常に難しく、ハードウェア要件が高く、データセットのダウンロードサイクルが長く、必要なスペースが大きいため、一般的な開発者は導入を躊躇します。そのため、オープンソースコミュニティは、AlphaFold2 の使用可能なバージョンを実現するために懸命に取り組んできました。

コロンビア大学のモハメッド・アルクライシ教授らが実装したOpenFoldの総トレーニング時間はA100で約10万時間だったが、約3,000時間で90%の精度を達成した。

OpenFold は、オリジナルの AlphaFold2 と同等か、それよりわずかに優れた精度を備えています。これは、OpenFold のトレーニングセットが少し大きいためと考えられます。

OpenFold の主な利点は、推論速度が大幅に向上することです。タンパク質配列が短い場合、OpenFold の推論速度は AlphaFold2 の 2 倍になります。さらに、カスタム CUDA カーネルを使用することで、OpenFold はより少ないメモリを使用してより長いタンパク質配列を推測できます。

OpenFold の紹介

OpenFold は、DeepMind 独自のアブレーションテストでパフォーマンスが低かった非推奨の「モデルアンサンブル」機能を除き、元のオープンソース推論コード (v2.0.1) のほぼすべての機能を再現します。

OpenFold は、DeepSpeed の有無にかかわらず、完全精度または bfloat16 でトレーニングできます。 AlphaFold2 の本来のパフォーマンスを実現するために、チームは OpenFold をゼロからトレーニングし、モデルの重みとトレーニングデータを公開しました。トレーニングデータには、約 400,000 個の MSA および PDB70 テンプレートファイルが含まれています。 OpenFold は、AlphaFold の公式パラメータを使用したタンパク質推論もサポートします。

他の実装と比較して、OpenFold には次の利点があります。

短いシーケンスの推論: 1500 アミノ酸残基未満のチェーンの GPU での推論を高速化します。
長いシーケンスの推論: この研究で実装された低メモリアテンションを使用して非常に長いチェーンを推論することで、OpenFold は単一の A100 で 4,000 残基を超えるシーケンス構造を予測でき、CPU オフロードによりさらに長いシーケンスを予測できます。
メモリ効率が高いトレーニングと推論中、FastFold カーネルから変更されたカスタム CUDA アテンションカーネルは、同等の FastFold および既存の PyTorch 実装と比較して、それぞれ 4 倍と 5 倍少ない GPU メモリを使用します。
効率的なアライメントスクリプト: チームは、オリジナルの AlphaFold HHblits/JackHMMER パイプラインまたは ColabFold と MMseqs2 を使用して、何百万ものアライメントを生成しました。

Linuxシステムでのインストールと使用

開発チームは、Miniconda をローカルにインストールし、conda 仮想環境を作成し、すべての Python 依存関係をインストールし、2 セットのモデルパラメーターを含む便利なリソースをダウンロードするスクリプトを提供します。

次のコマンドを実行します。

 スクリプト/ install_third_party_dependencies.sh

次のコマンドを使用して環境をアクティブ化します。

 ソーススクリプト/ activate_conda_env.sh

非アクティブ化コマンド:

 ソーススクリプト/ deactivate_conda_env.sh

アクティブ化された環境でOpenFoldのCUDAカーネルをコンパイルする

 python3 セットアップ.py インストール

/usr/bin に HH-suite をインストールします。

 # スクリプト/install_hh_suite.sh

OpenFold と AlphaFold のトレーニングに使用するデータベースをダウンロードするには、次のコマンドを使用します。

 bash スクリプト/ download_data . sh データ/

DeepMind の事前トレーニング済みパラメータのセットを使用して 1 つ以上のシーケンスに対して推論を実行する場合は、次のコードを実行できます。

 python3 run_pretrained_openfold.py \
 ファスタディレクトリ\
 データ/ pdb_mmcif / mmcif_files / \
 --uniref90_database_path データ/ uniref90 / uniref90.fasta \
 -- mgnify_database_path データ/ mgnify / mgy_clusters_2018_12 . fa \
 --pdb70_database_path データ/ pdb70 / pdb70 \
 --uniclust30_database_path データ/ uniclust30 / uniclust30_2018_08 / uniclust30_2018_08 \
 --出力ディレクトリ./ \
 --bfd_database_path データ/ bfd / bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
 --model_device "cuda:0" \
 --jackhmmer_binary_path lib / conda / envs / openfold_venv / bin / jackhmmer \
 --hhblits_binary_path lib / conda / envs / openfold_venv / bin / hhblits \
 --hhsearch_binary_path lib / conda / envs / openfold_venv / bin / hhsearch \
 --kalign_binary_path lib / conda / envs / openfold_venv / bin / kalign
 --config_preset "モデル_1_ptm "
 --openfold_checkpoint_path openfold / リソース/ openfold_params / finetuning_2_ptm . pt