ご注文はリード化合物ですか?〜医薬化学録にわ〜

自分の勉強や備忘録などを兼ねて好き勝手なことを書いていくブログです。

とりあえず PyTorch: 最低限これで動く

ニューラルネットワーク(深層学習)ブームは 2022 年現在も続いており、毎日新しい研究がニュースや SNS で話題となっています。ケモインフォマティクス分野でも例外ではなく、QSAR モデルや構造生成器、逆合成解析を始めとして様々な深層学習を用いた研究がされています。
深層学習のライブラリは Tensorflow や Caffe なども知られていますが、この記事を執筆している時点では、PyTorch が盛んに使われています。しかし、PyTorch でとりあえず QSAR モデルを組みたい、という場合に、意外と日本語のページがヒットしないような気がします。
そこで、非常に最低限ですが、とりあえず回帰モデルを組みたい、という場合に使える基本的なコードを作成してみました。

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
from torch.optim import Adam
from sklearn.metrics import r2_score

# データの読み込み
train_X = pd.read_csv("../data/LogS_train_scaled.csv", index_col = 0)
test_X = pd.read_csv("../data/LogS_test_scaled.csv", index_col = 0)

train_y = pd.read_csv("../data/LogS_train_value.csv", index_col = 0)
test_y = pd.read_csv("../data/LogS_test_value.csv", index_col = 0)

# データを Tensor 型にする
train_X_tensor = torch.Tensor(np.array(train_X)).float()
train_y_tensor = torch.Tensor(np.array(train_y)).float()
train_tensor = TensorDataset(train_X_tensor, train_y_tensor)

test_X_tensor = torch.Tensor(np.array(test_X)).float()
test_y_tensor = torch.Tensor(np.array(test_y)).float()
test_tensor = TensorDataset(test_X_tensor, test_y_tensor)

# Data loader を作成 
train_loader = DataLoader(train_tensor, batch_size = 32, shuffle = True)
test_loader = DataLoader(test_tensor, batch_size = 32, shuffle = False)

# モデルのクラスを設定
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_1 = nn.Linear(94, 32)
        self.linear_2 = nn.Linear(32, 1)
        
        self.norm = nn.BatchNorm1d(32)
        
    def forward(self, x):
        x = F.relu(self.linear_1(x))
        x = self.norm(x)
        x = self.linear_2(x)
        
        return x

# モデルのインスタンスを作成
model = Model()

# 最適化関数の設定
optimizer = Adam(model.parameters())

# 損失関数の設定
loss_function = nn.MSELoss()

# モデルの学習
epoch = 20
model.train()
for i in list(range(epoch)):
    
    epoch_loss = 0
    
    # バッチごとの処理
    for batch_X, batch_y in train_loader:
        
        # トレーニングデータに対する誤差の算出
        batch_pred_y = model(batch_X)
        batch_loss = loss_function(batch_pred_y, batch_y)
        
        # 誤差逆伝播と重み更新
        optimizer.zero_grad()
        batch_loss.backward()
        optimizer.step()
        
        epoch_loss += batch_loss.item()
        
    # エポックごとの誤差を表示
    print(epoch_loss)

# モデルの推論
model.eval()
with torch.no_grad():
    test_loss = 0
    pred_y = torch.Tensor()
    
    # トレーニングと同様の処理を行うが、重みの更新はしない
    for batch_X, batch_y in test_loader:
        
        batch_pred_y = model(batch_X)
        batch_loss = loss_function(batch_pred_y, batch_y)        
        test_loss += batch_loss.item()
        pred_y = torch.cat([pred_y, batch_pred_y])
        
    # 誤差を表示
    print(test_loss)

# Numpy 型に戻した後に R2 を算出
pred_y = pred_y.detach().numpy()
print(r2_score(test_y, pred_y))

# 構築したモデルの保存
torch.save(model.state_dict(), "../model/nn_model.pth")

今回は、94 次元の記述子から、pLogS(連続値)を予測するモデルとなっています。データサイズ はトレーニングデータ 1,200、テストデータ 94 なのでとても小さいです。大体の目安ですが、10,000 データ以上の場合、深層学習は力を発揮します。
最初のポイントは、データを Tensor 型にするところです。PyTorch は基本的にデータを Tensor 型にしないと計算してくれないので、Numpy のフレームから変更する必要があります。また、PyTorch では int 型や float64 型などは学習に用いることができないので、float() で float32 型にしてあげます。説明変数(X)と目的変数(y)を一つの DataFrame にまとめてから、バッチを作成する関数である DataLoader に入力します。ここでバッチサイズを決めます。バッチサイズは大雑把ですが、データ数にルートを取った値ぐらいを試すのがちょうど良い気がします(データ数が数マンだったら 256 ぐらい、といった感じ)。

深層学習モデルのクラスは、 __init__ と forward から構成されます。__init__ にモデルで使う全結合層や畳み込み層などを定義します。__init__ に super などがついている場合もありますが、少し古い書き方だそうで、基本的に今回の例のようなシンプルな書き方で問題ありません。forward で関数を組み合わせて、モデルの形状を定義します。今回は batch normalization と活性化関数(ReLU)を挟んだ、二層の全結合層モデルとなっています。データサイズが少ない(20,000 以下ぐらい?)場合は、モデルを複雑にしても予測精度は向上しない気がします。今回は回帰問題なので出力層はそのままですが、判別問題の場合は、各クラスの確率を出力したいので、ソフトマックス関数を追加することを忘れないようにしましょう。

モデルのインスタンスと最適化関数、損失関数を定義します。最適化関数は SGD でも良いのですが、とりあえずよく使われている Adam を今回は採用します。model.parameters() はおまじないとして、入力を忘れないようにして下さい。学習率(lr)は、データサイズが大きい場合に小さくすると良いです。1e-4 をとりあえず試して見るのが良いかもしれません。損失関数は平均二乗誤差を用います。判別問題の場合はクロスエントロピーを用いて下さい。

いよいよモデルの学習です。エポック数は本当は early stopping を実装して決めるべきなのですが、今回は適当に 20 に設定します。定義した DataLoader から for 文でデータを取り出します。モデルのインスタンスにデータを入力して予測値を算出し、誤差を算出します。その後に、誤差逆伝播法を実行します。勾配の初期化、誤差逆伝播の具体的な計算、最適化関数のパラメータ更新などを行っているそうですが、ここもまずはおまじないということで良いでしょう。

モデルの学習が出来たら、推論を行います。折角学習した重みに変更を加えたくないので、model.eval() と torch.no_grad() を両方実行します。計算方法は学習時と基本的に同じですが、誤差逆伝播法は当然行いません。予測値を、torch.cat を使って保存していきます。

出力値の R2 を計算してみます。Tensor 型から NumPy 型に変換しますが、勾配の学習に用いた Tensor は、一度 detach() を挟む必要があります。とりあえずつけてから numpy() を実行するようにしておくとエラーが起きにくいと思います。

シード値をつけていないので、多少ズレはあるのですが、R2 は大体 0.88 前後と、適当に組んだ割にはまあまあ良い値となりました。

保存したモデルの読み込みは以下の通りです。

# モデルの読み込み
loaded_model = Model()
loaded_model.load_state_dict(torch.load("../model/nn_model.pth"))

モデル構築に用いたクラスを読み込み、そこに学習したモデルの重みを入れるイメージです。

今回のコードは CPU で実行しました。GPU に対応させる場合は、バッチのデータとモデルに cpu() を追加して下さい。また、用いる GPU のバージョンに対応した PyTorch が適切にインストールされているかも注意点です。

最低限の内容ですが、とりあえず動かしたい、という場合の参考になれば幸いです。