PyTorchベースの計算化学用ディープラーニングライブラリOpenChem

#はじめに
化合物でディープラーニングといえば、DeepChem (https://deepchem.io/) や　Chainer Chemistory(https://chainer-chemistry.readthedocs.io/en/latest/) が有名だが、今回OpenChemというライブラリを紹介したい。

#OpenChemとは
計算化学およびドラッグデザイン研究のためのディープラーニングツールキットである。PyTorchをバックエンドに使用しており、ライセンスはMITライセンスである。

#どんな特徴があるの？
その特徴として以下があげられる

• モジュラー設計により、設定ファイルのみでモジュールを簡単に組み合わせて構築することができる。
• マルチGPUをサポートしており、高速に学習することができる。
• データ前処理のためのユーティリティ等も備えている。
• テンソルボードという可視化ツールにも対応している。

モジュラー設計について簡単に説明すると、DeepChemで新しいモデルを作成する場合、入力から出力までほぼ１から開発する必要があるのに対し、OpenChemではエンコーダー、デコーダー、埋め込みレイヤーなどの標準的なディープニューラルネットワークのブロックが用意されており、これらを組み合わせることとで既存のものを再利用しながら新しいモデルを比較的容易に開発することができるのである。

#どんなタイプのモデルが作れるの？

• 分類（2クラス、多クラス)
• 回帰
• マルチタスク
• 生成モデル

#どんなデータを扱えるの？

• SMILES文字列やアミノ酸配列などの文字の配列。
• 分子グラフ（SMILES文字列から生成）

どんなモジュールが用意されているの？

• トークンの埋め込み
• リカレントニューラルネットワークエンコーダー
• グラフ畳み込みニューラルネットワークエンコーダー
• 多層パーセプトロン

#インストール
以下の環境でインストールしてみる。

環境

• Windows 10
• python 3.7
• rdkit-2020.09.1
• pytorch 1.7.0
• scikit-learn 0.23.2
• tensorboard 2.4.0

手順

以下のコマンドを順に実行すればよい。
ちなみに、tqdm以降はドキュメントに書かれていなかったが、試行錯誤しているうちに必要だと分かったため追加している。

$ conda create -n openchem python=3.7
$ conda activate openchem
$ git clone https://github.com/Mariewelt/OpenChem.git
$ cd OpenChem
$ conda install --yes --file requirements.txt
$ conda install -c rdkit rdkit nox cairo
$ conda install pytorch torchvision -c pytorch
$ conda install tqdm
$ conda install tensorboard
$ conda install networkx

#モデルの定義と学習はどうやるの？

OpenChemのモデルは、設定ファイル（Pythonのディクショナリ形式でパラメータを指定）によって定義することができる。
設定ファイルには、以下を含める必要がある。

• モデルの実行/トレーニング/評価の方法を定義するパラメータ
• モデルアーキテクチャを定義するパラメータ

また、モデルの作成と分散プロセスの起動を処理する2つのPythonファイルである「run.py」と「launch.py」も必要となる。

##runc.pyの引数

項目	説明
nproc_per_node	ノードあたりのプロセス数。 ノード上のGPUの数と同じである必要がある。

##lanch.pyの引数

項目	説明
config_file	モデルが定義されているPython設定ファイルへのパス。
mode	train、train_evalまたはeval。
continue_learning	この引数が指定されている場合、トレーニングは最新のチェックポイントから再開される

#設定ファイル
設定ファイルにはモデルが含まれている必要がある。
モデルは、OpenChemModelおよびディクショナリmodel_paramsから派生したクラスである必要がある。
以下は、モデルアーキテクチャに関連しないすべてのモデルに共通のパラメータの説明である。

項目	説明
task	モデルによって解決されるタスクを指定する。 classification, regression , multitaskのいずれか
train_data_layer	トレーニングデータ用のpytorchデータセット。 --mode = evalの場合、Noneになる。 OpenChemは現在、SMILES、Graph、およびMoleculeProteinデータセットを作成するためのユーティリティを提供していている。
val_data_layer	検証データ用のpytorchデータセット。 --mode = trainの場合、Noneになる。
print_every	ログ出力の頻度
save_every	モデルがチェックポイントに保存される頻度。
logdir	モデルチェックポイントとテンソルボードログが保存されるフォルダーへのパス。
use_clip_grad	勾配クリッピングを使用するかどうか。
max_grad_norm	勾配クリッピングが使用されている場合のパラメーターの最大ノルム。
batch_size	バッチサイズ
num_epochs	エポック数。 --mode = evalの場合、Noneになる
eval_metrics	評価のためのユーザー定義関数、指標。 --mode = trainの場合、Noneになる。
criterion	pytorchの損失、モデルの損失
optimizer	pytorchオプティマイザー、モデルをトレーニングするためのオプティマイザー。 --mode = evalの場合、Noneになる。

その他のパラメータは、モデルアーキテクチャに固有のものとなるので、APIドキュメントやその他のチュートリアルを確認してほしい。

##ジョブの起動

以下4つのGPUを備えたノードで実行されるジョブの例である。1つのGPUで単一のプロセスを実行することもできる。

python launch.py --nproc_per_node=4 run.py --config_file="./my_config.py" --mode="train"

#チュートリアルをやってみよう
分子フィンガープリントからlogP値を予測するための単純な多層パーセプトロンニューラルネットワークを構築することにより、OpenChemでのモデル構築の基本について試してみる。

データのロード

まず、ファイルからデータを読み取る必要がある。
サンンプルのデータは、PyTorchのリポジトリの ./benchmark_datasets/logp_dataset/logP_labels.csv にあるものを使う。

OpenChemは、複数の列を持つテキストファイルを処理することができる。
read_smiles_property_file では読み取る列、区切り文字を指定することができる。
読み取る列（cols_to_read）では、最初にSMILESの列、次の列にはラベルを指定する。

import numpy as np 
from openchem.data.utils import read_smiles_property_file
data = read_smiles_property_file('./benchmark_datasets/logp_dataset/logP_labels.csv',
                                 delimiter=",",
                                 cols_to_read=[1, 2],
                                 keep_header=False)

read_smiles_property_fileの戻り値は、ファイルから読み取られた列と同じ数のオブジェクトを含むリストであり、上の例では、data[0]にはSMILESが含まれ、残りはすべてラベルとなっている。

smiles = data[0]
labels = np.array(data[1:])
labels = labels.T

データを読み取った後、scikit-learnの関数を使用してトレーニングセットとテストセットに分割し、新規にファイルに保存する。

from sklearn.model_selection import train_test_split
from openchem.data.utils import save_smiles_property_file

X_train, X_test, y_train, y_test = train_test_split(smiles, labels, test_size=0.2,
                                                    random_state=42)
save_smiles_property_file('./benchmark_datasets/logp_dataset/train.smi', X_train, y_train)
save_smiles_property_file('./benchmark_datasets/logp_dataset/test.smi', X_test, y_test)

PyTorchデータセットの作成

続いてPyTorchで処理するためのデータセットを作成する。OpenChemには、データ型に基づいてPyTorchデータセットを作成するための複数のユーティリティがある。
今回、SMILES文字列をユーザー定義関数を使用して特徴のベクトルに変換するFeatureDatasetを使用する。

from openchem.data.feature_data_layer import FeatureDataset
from openchem.data.utils import get_fp


train_dataset = FeatureDataset(filename='./benchmark_datasets/logp_dataset/train.smi',
                               delimiter=',', cols_to_read=[0, 1],
                               get_features=get_fp, get_features_args={"n_bits": 2048})
test_dataset = FeatureDataset(filename='./benchmark_datasets/logp_dataset/test.smi',
                              delimiter=',', cols_to_read=[0, 1],
                              get_features=get_fp, get_features_args={"n_bits": 2048})
predict_dataset = FeatureDataset(filename='./benchmark_datasets/logp_dataset/test.smi',
                                delimiter=',', cols_to_read=[0],
                                get_features=get_fp, get_features_args={"n_bits": 2048},
                                return_smiles=True)

ユーザー定義関数は、FeatureDatasetの引数get_featuresに渡され、さらにユーザー定義関数の引数はget_features_argsにディクショナリとして渡すことができる。

この例では、ユーザー定義関数として、RDKitフィンガープリントを生成するopenchem.data.utils.get_fp を利用している。
この関数はn_bitsにフィンガープリントのビット数を指定することができる。

read_smiles_property_finction と同様、Datasetには cols_to_read およびdelimiter引数を指定する。

最終的に、train_dataset、test_dataset、predict_dataset の3つのデータセットを作成している。train_datasetとtest_datasetは、それぞれトレーニングと評価に使用する。
これらのデータセットでは、cols_to_read にSMILES文字列だけでなく、ラベルを持つ列のインデックスも含める必要がある。

トレーニングが完了した後、predict_datasets を使用して、新しいサンプルの予測結果を取得する。
predict_datasets ではラベルは必要ないため、cols_to_read引数には、SMILES文字列に対する列のインデックスのみ指定すればよい。また、predict_datasetには、引数return_smiles = True を指定する。

OpenChemモデルの作成とパラメータの定義

続いて、モデルタイプとモデルパラメータを指定する。ここでは多層パーセプトロンモデルであるMLP2Labelモデルを使用しており、特徴ベクトルからラベルを予測する。

from openchem.models.MLP2Label import MLP2Label
model = MLP2Label

model_params = {
    'task': 'regression',
    'random_seed': 42,
    'batch_size': 256,
    'num_epochs': 101,
    'logdir': 'logs/logp_mlp_logs',
    'print_every': 20,
    'save_every': 5,
    'train_data_layer': train_dataset,
    'val_data_layer': test_dataset,
    'predict_data_layer': predict_dataset,
    'eval_metrics': r2_score,
    'criterion': nn.MSELoss(),
    'optimizer': Adam,
    'optimizer_params': {
        'lr': 0.001,
    },
    'lr_scheduler': StepLR,
    'lr_scheduler_params': {
        'step_size': 15,
        'gamma': 0.9
    },
    'mlp': OpenChemMLP,
    'mlp_params': {
        'input_size': 2048,
        'n_layers': 4,
        'hidden_size': [1024, 512, 128, 1],
        'dropout': 0.5,
        'activation': [F.relu, F.relu, F.relu, identity]
    }

モデルの学習

上記で記載した「データのロード」から「OpenChemモデルの作成とパラメータの定義」までに記載したPythonのコードをファイルに保存する。ここではexample_configsフォルダーにあるlogp_mlp_config.pyファイルに同じものが保存されているので、それをそのまま使うことにする。

そして、コマンドラインから次のコマンドを実行して学習を開始することができる。

python launch.py --nproc_per_node=1 run.py --config_file=example_configs/getting_started.py  --mode="train_eval"

うまくいけば、モデル構成、全体的なトレーニングの進行状況、トレーニングの損失、検証の損失、およびR^2スコアである検証メトリックが出力される。

トレーニング済みモデルを予測モードでさらに実行して、新しいサンプルの予測結果を得るには、コマンドラインから次のコマンドを実行する必要がある。

python launch.py --nproc_per_node=1 run.py --config_file=example_configs/getting_started.py  --mode="predict"

予測結果

こんな感じで得られる。

less logs/logp_mlp_logs/predictions.txt

CCCCCCN(C)N=Nc1ccc(C(N)=O)cc1,4.7723255
NC1C2CCC1c1ccccc12,2.6808066
Nc1nccc2ccccc12,1.7958283
CC(C)(O)CCS(=O)(=O)c1ccc(S(N)(=O)=O)cc1,0.10045385
CCCNC(=O)c1ccc[nH]1,0.9938375
O=C(Nc1ccc(Cl)cc1)c1cccc([N+](=O)[O-])c1O,4.165857
COC(=O)C(C#N)=NNc1ccc2c(c1)OC(F)(F)C(F)(F)O2,4.892629

#おわりに
この後は、Tox21のデータを使ったチュートリアル等でさらに理解を深めるとよい。また、プログラミング力は必要になるが、試行錯誤しながら論文のモデルを実装してみると面白い。DeepChemに比べるとカスタマイズできる感じが半端なくあるし、バックエンドもPyTorchなのでソースも理解しやすい。まさに、化合物によるDeepLearningの入門にうってつけのライブラリといえるだろう。

#URL

• https://mariewelt.github.io/OpenChem/html/index.html