化合物VAEを作ろうとしたのですが、VAEは思ったよりややこしかったのでとりあえず化合物AEをしました。
はじめに
化合物構造（smiles）を入力し、次元圧縮した後に化合物構造(smiles)を復元するautoencoderモデルを作成します。
VAEではないので化合物生成には使えませんが、エンコードされた層（図中のcompressed representation）が分子フィンガープリントとして使えます。
※図は
Building Autoencoders in Kerasより改変して作成
データセット
Maximum Unbiased Validation (MUV) Data Set（https://pubs.acs.org/doi/10.1021/ci8002649）を使用します。MUVはvirtual screeningのベンチマーク用に設計された9万ほどの多様な低分子化合物のデータセットです。
deepchem/datasets at master · deepchem/deepchem · GitHub
deepchem/datasets/にmuv.csv.gzとしてcsv形式で置いてあります。
データセットの前処理
前処理としてデータセットをone-hot encodingします。
#データセットのロード import pandas as pd dataset = pd.read_csv("data/muv.csv") smiles = dataset["smiles"].values # データセットに含まれる語彙を収集 chars = set() for smi in smiles: chars = chars.union(set(char for char in smi)) chars = sorted(list(chars)) #語彙辞書の作成 char_dict = dict((char, i+1) for i,char in enumerate(chars)) char_dict["<EOS>"] = 0 num_dict = dict((i+1,char) for i,char in enumerate(chars)) num_dict[0]="<EOS>" chars_size=len(char_dict) max_length = max(len(s) for s in smiles)
>>>print(char_dict) {'#': 1, '(': 2, ')': 3, '+': 4, '-': 5, '/': 6, '1': 7, '2': 8, '3': 9, '4': 10, '5': 11, '6': 12, '=': 13, 'B': 14, 'C': 15, 'F': 16, 'H': 17, 'N': 18, 'O': 19, 'S': 20, '[': 21, '\\': 22, ']': 23, 'c': 24, 'l': 25, 'n': 26, 'o': 27, 'r': 28, 's': 29, '<EOS>': 0} >>>print(num_dict) {1: '#', 2: '(', 3: ')', 4: '+', 5: '-', 6: '/', 7: '1', 8: '2', 9: '3', 10: '4', 11: '5', 12: '6', 13: '=', 14: 'B', 15: 'C', 16: 'F', 17: 'H', 18: 'N', 19: 'O', 20: 'S', 21: '[', 22: '\\', 23: ']', 24: 'c', 25: 'l', 26: 'n', 27: 'o', 28: 'r', 29: 's', 0: '<EOS>'} >>>print(len(char_dict)) 30 >>>print(max_length) 86
max_lengthが86であることから、データセット内で最長のsmilesは86文字であることがわかります。
またデータセット内のsmilesは29種類の文字+<EOS>で表現できることがわかります。
※<EOS>は列の終わりを示すため追加
# One-Hot処理 X = np.zeros( (len(smiles), max_length+1, chars_size),dtype=np.bool) for n, smi in enumerate(smiles): for m, char in enumerate(smi): char_idx = char_dict[char] X[n,m,char_idx] = 1 else: X[n,len(smi),0]=1 continue
>>>print(X.shape) (93127, 87, 30)
one-hot処理によりすべてのサンプルが87（文字数）×30（語彙数）の行列になりました。
max_length+1としているのは列の終わりを示す<EOS>を追加したからです。
モデルの定義
from keras.layers import Input, Dense, GRU, RepeatVector, TimeDistributed, Conv1D, Flatten from keras.models import Model def build_autoencoder(max_length,chars_size,pretrained_weight=None): #encoder shape=(max_length,chars_size) encoder_inputs = Input(shape=shape) conv1 = Conv1D(filters=27, kernel_size=9, strides=1, padding='same',activation='relu')(encoder_inputs) conv2 = Conv1D(filters=18, kernel_size=9, strides=1, padding='same',activation='relu')(conv1) conv3 = Conv1D(filters=9, kernel_size=9, strides=1, padding='same',activation='relu')(conv2) flat = Flatten()(conv3) dense = Dense(512)(flat) encoded = dense #decoder decoder_inputs = RepeatVector(shape[0])(encoded) gru1 = GRU(512, return_sequences=True)(decoder_inputs) gru2 = GRU(256, return_sequences=True)(gru1) gru3 = GRU(125, return_sequences=True)(gru2) dense = TimeDistributed(Dense(chars_size, activation='softmax'))(gru3) decoder_outputs = dense model = Model(encoder_inputs, decoder_outputs) if pretrained_weight: model.load_weights(pretrained_weight) return model
>>>model = build_autoencoder(max_length+1, chars_size, pretrained_weight=None) >>>model.summary() _________________________________________________________________ Layer (type) Output Shape Param # ================================================================= input_21 (InputLayer) (None, 87, 30) 0 _________________________________________________________________ conv1d_58 (Conv1D) (None, 87, 27) 7317 _________________________________________________________________ conv1d_59 (Conv1D) (None, 87, 18) 4392 _________________________________________________________________ conv1d_60 (Conv1D) (None, 87, 9) 1467 _________________________________________________________________ flatten_16 (Flatten) (None, 783) 0 _________________________________________________________________ dense_27 (Dense) (None, 512) 401408 _________________________________________________________________ repeat_vector_13 (RepeatVect (None, 87, 512) 0 _________________________________________________________________ gru_39 (GRU) (None, 87, 512) 1574400 _________________________________________________________________ gru_40 (GRU) (None, 87, 256) 590592 _________________________________________________________________ gru_41 (GRU) (None, 87, 125) 143250 _________________________________________________________________ time_distributed_13 (TimeDis (None, 87, 30) 3780 ================================================================= Total params: 2,726,606 Trainable params: 2,726,606 Non-trainable params: 0 _________________________________________________________________
DecoderではGRU層を3層重ねていますがユニット数は徐々に減少させていきます。これはRNN層が遅いので計算速度を確保するためです。 なぜかは知りませんが化合物のモデルではLSTMではなくGRUがよく使われるようです。
DeepChemのSeqToSeqモデル
https://deepchem.io/docs/_modules/deepchem/models/tensorgraph/models/seqtoseq.html
でもGRUが採用されています。
トレーニング
from keras.callbacks import * X_train = X[:70000, :, :] X_test = X[70000:, :, :] model = build_autoencoder(max_length+1,chars_size,pretrained_weight=None) model.compile(optimizer='adam', loss='categorical_crossentropy') callbacks = [] callbacks.append(ModelCheckpoint(filepath="checkpoints/AE_{epoch:02d}.hdf5",save_weights_only=True)) history = model.fit(X_train, X_train, epochs=20, batch_size=512, shuffle=True, callbacks=callbacks, validation_data=(X_test, X_test))
もっと長く学習させれば精度向上しそうな感じですが面倒なので打ち切りました。
検証
One-hot表現をSmiles文字列に復元する関数を作成し、モデルの検証を行います。
def Onehot2smiles(sample): sample_onehot = sample.reshape(1,87,30) sample_onehot_pred = model.predict(sample_onehot) sample_truth = [] for idx in range(max_length): char_idx = np.where(sample_onehot[0,idx,:]==1)[0] if not list(char_idx): break char_idx = char_idx[0] char = num_dict[char_idx] sample_truth.append(char) truth = "".join(sample_truth) sample_pred = [] for idx in range(max_length): char_idx = np.where(sample_onehot_pred[0,idx,:]==max(sample_onehot_pred[0,idx,:]))[0] if not list(char_idx): break char_idx = char_idx[0] char = num_dict[char_idx] sample_pred.append(char) pred = "".join(sample_pred) print("Input:",truth) print("Output:",pred) print()
>>>for i in range(15): >>> idx = np.random.randint(X_test.shape[0]) >>> Onehot2smiles(X_test[idx]) Input: CCN1C(=O)C(CC(=O)Nc2ccc(OC)cc2)N(Cc2cccnc2)C1=S<EOS> Output: CCN1C(=O)C(CC(=O)Nc2ccc(OC)cc2)N(Cc2cccnc2)C1=S<EOS> Input: Cc1cc(=O)oc2cc(OC(C)C(=O)O)c(Cl)cc12<EOS> Output: Cc1cc(=O)oc2cc(OC(C)C(=O)O)c(Cl)cc12<EOS> Input: CNC(=O)c1c(N)sc2c1CCCCC2<EOS> Output: CNC(=O)c1c(N)sc2c1CCCCC2<EOS> Input: COc1ccc(N2CCN(Cc3cc(=O)oc4cc(-c5ccccc5)c(O)cc34)CC2)cc1<EOS> Output: COc1ccc(N2CCN(Cc3cc(=O)nc4cc(-c5ccccc4)c(C)ccc3)CC2)cc1<EOS> Input: CC(C)(C)OC(=O)N1CC2Cc3[nH]c(=O)ccc3C(C2)C1<EOS> Output: CC(C)(C)OC(=O)N1CCN(nnnnH]c(=O)ccc2C#C))C1<EOS> Input: O=C(Nc1ccc(Cl)c(S(=O)(=O)N2CCOCC2)c1)c1ccccn1<EOS> Output: O=C(Nc1ccc(Cl)c(S(=O)(=O)N2CCOCC2)c1)c1ccccn1<EOS> Input: CCc1nc2c3cnn(-c4cc(Cl)ccc4C)c3ncn2n1<EOS> Output: CCc1nc2c3c(c(-c4cc(Cl)ccc3C)c3cnc2n1<EOS> Input: Cc1cccc(Nc2ccccc2C(=O)NCc2ccco2)c1C<EOS> Output: Cc1cccc(Nc2ccccc2C(=O)NCc2ccco2)c1C<EOS> Input: O=C(c1cccc(S(=O)(=O)N2CCCC2)c1)c1cccc(S(=O)(=O)N2CCCC2)c1<EOS> Output: O=C(c1cccc(S(=O)(=O)N2CCCC2)c1)c1cccc(S(=O)(=O)N2CCOC2)cc1<EOS> Input: C=CCc1cccc2cc(C(=O)NCc3cccnc3)c(=O)oc12<EOS> Output: C=CCc1cccc2cc(C(=O)NCc3ccccc3)c(=O)oc12<EOS> Input: COc1ccc(C(=O)Nc2nnc(Cc3cccs3)s2)cc1OC<EOS> Output: COc1ccc(C(=O)Nc2nnc(Cc3cccs3)s2)cc1OC<EOS> Input: CCC1CCCCN1CCCNC(=O)Cn1nc(-c2ccc(C)cc2)ccc1=O<EOS> Output: CCC1CCCCN1CCNNC(=O)Cn1nc(-c2ccc(C)cc2)ccc1=O<EOS> Input: O=C(Cn1[nH]c(=O)c2ccccc2c1=O)OCC(=O)N(Cc1ccccc1)Cc1ccccc1<EOS> Output: O=C(Cn1[nH]c(=O)c2ccccc2c1=O)OCC(=O)N(Cc1ccccc1)cc1cccc1<EOS> Input: CN1CCN(c2nc3c(c(=O)n(C)c(=O)n3C)n2Cc2cccc3ccccc23)CC1<EOS> Output: CN1CCN(c2nc3c(c(=O)n(C)c(=O)n3C)n2Cc2cccc3ccccc23)CC1<EOS> Input: CCCc1nc(SC2CCCC2)c2c(=O)n(C)c(=O)n(C)c2n1<EOS> Output: CCCc1nc(SC2CCCC2)c2c(=O)n(C)c(=O)n(C)c2n1<EOS>
それなりに良い感じでしょうか？
おわりに
kerasでのRNNの日本語での解説ってあんまりないですね。TimeDistributedレイヤーとか理解するまでいろいろ調べました。
付録：試行錯誤の経過
lossが0.3あたりからまともにsmilesが復元するようになってきた印象です。
- encoder: LSTM decoder:LSTMのモデルを試すも全くうまくいかず (loss>1.3)
- データの表現を
[C,N,=,C,=,O,<pad>,<pad>,<pad>]という形式から[C, N, =, C, =, O, <EOS>, 0, 0]に変更 (loss>1.0) - encoder: GRU decoder:GRUに変更（精度向上せず）
- encoder: Conv1D decoder:GRUのモデルに変更
参考
Building Autoencoders in Keras
A ten-minute introduction to sequence-to-sequence learning in Keras
