CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: TabNet を使ってみる

NeuralNetwork PyTorch LightGBM 機械学習 scikit-learn Seaborn Ubuntu20.04LTS Linux

一般に、テーブルデータの教師あり学習では、勾配ブースティング決定木の性能の良さについて語られることが多い。 これは、汎化性能の高さや前処理の容易さ、学習・推論の速さ、解釈可能性の高さなどが理由として挙げられる。 一方で、ニューラルネットワークをテーブルデータに適用する取り組みについても、以前から様々な試みがある。 今回は、その中でも近年機械学習コンテストにおいて結果を残している TabNet というモデルを試してみる。 TabNet には Unsupervised pre-training と Supervised fine-tuning を組み合わせた学習方法や、モデルの解釈可能性を向上させる試みなどに特徴がある。

使った環境は次のとおり。

$ cat /etc/lsb-release 
DISTRIB_ID=Ubuntu
DISTRIB_RELEASE=20.04
DISTRIB_CODENAME=focal
DISTRIB_DESCRIPTION="Ubuntu 20.04.5 LTS"
$ uname -srm
Linux 5.15.0-56-generic x86_64
$ python -V
Python 3.10.9
$ pip list | grep -i torch
pytorch-tabnet               4.0
torch                        1.13.1

もくじ

下準備

あらかじめ利用するパッケージをインストールしておく。 pytorch-tabnet は非公式ながら PyTorch を使った TabNet の最も有名な実装となる。

$ pip install pytorch-tabnet seaborn scikit-learn

回帰タスク (Diamonds データセット)

今回は試しに Diamonds データセット 1 について、価格 (price) カラムを目的変数として RMSLE で評価してみよう。

早速だけど以下にサンプルコードを示す。 TabNet の特徴として Unsupervised pre-training で学習した重みを Supervised fine-tuning の学習に転用する点が挙げられる。 これらは、それぞれ TabNetPretrainer と TabNetRegressor というクラスとして表現されている。 いずれも scikit-learn の Transformer API を実装しているため簡便に扱うことができる。 なお、Unsupervised pre-training は必須ではないことから、精度を落としても構わなければ省略できる。 ハイパーパラメータなどについては、デフォルトからいじっているのは Entity Embedding の埋め込み次元数と学習率、エポック数くらい。

import math

import numpy as np
import pandas as pd
import seaborn as sns
from pytorch_tabnet.pretraining import TabNetPretrainer
from pytorch_tabnet.tab_model import TabNetRegressor
from sklearn.compose import ColumnTransformer
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import OrdinalEncoder
import torch
from torch import nn
from matplotlib import pyplot as plt


def rmsle_metric(y_pred, y_true):
    """Root Mean Squared Logarithmic Error を計算する関数"""
    clamped_y_pred = np.clip(y_pred, a_min=0., a_max=None)
    log_y_pred = np.log1p(clamped_y_pred)
    log_y_true = np.log1p(y_true)
    return mean_squared_error(log_y_pred, log_y_true, squared=False)


class RMSLELoss(nn.Module):
    """PyTorch で RMSLE を計算するモジュール"""

    def __init__(self):
        super().__init__()
        self.mse_loss = nn.MSELoss()

    def forward(self, y_pred, y_true):
        # 入力が -1 以下になると NaN になるため値に下限を設ける
        clamped_y_pred = torch.clamp(y_pred, min=0.)
        log_y_pred = torch.log1p(clamped_y_pred)
        log_y_true = torch.log1p(y_true)
        msle = self.mse_loss(log_y_pred, log_y_true)
        rmsle_loss = torch.sqrt(msle)
        return rmsle_loss


def main():
    # Diamonds データセットを読み込む
    df = sns.load_dataset('diamonds')
    x, y = df.drop(['price'], axis=1), df['price'].values

    # カテゴリ変数は Ordinal Encode する
    categorical_cut_pipeline = Pipeline(steps=[
        ("cut-ordinal", OrdinalEncoder())
    ])
    categorical_color_pipeline = Pipeline(steps=[
        ("color-ordinal", OrdinalEncoder())
    ])
    categorical_clarity_pipeline = Pipeline(steps=[
        ("clarity-ordinal", OrdinalEncoder())
    ])
    # 連続変数は標準化する
    numerical_pipeline = Pipeline(steps=[
        ("numerical", StandardScaler()),
    ])
    pipeline = ColumnTransformer(transformers=[
        ("categorical-cut-pipeline", categorical_cut_pipeline, ["cut"]),
        ("categorical-color-pipeline", categorical_color_pipeline, ["color"]),
        ("categorical-clarity-pipeline", categorical_clarity_pipeline, ["clarity"]),
        ("numerical-pipeline", numerical_pipeline, ["carat", "depth", "table", "x", "y", "z"])
    ])
    transformed_x = pipeline.fit_transform(x)

    # 最終的な汎化性能の評価に使うホールドアウト
    train_x, test_x, train_y, test_y = train_test_split(
        transformed_x,
        y.reshape(-1, 1),
        test_size=0.33,
        random_state=42,
        shuffle=True,
    )

    # 学習に使うために学習データと評価データを分割する
    train_x, eval_x, train_y, eval_y = train_test_split(
        train_x,
        train_y,
        test_size=0.2,
        random_state=42,
        shuffle=True,
    )

    # カテゴリ変数のインデックスとユニークな値の数を計算しておく
    categorical_index_with_modalities = {
        i: np.unique(transformed_x[:, i]).shape[0]
        for i in [0, 1, 2]
    }

    # 学習パラメータ
    tabnet_params = {
        # カテゴリ変数は Entity Embedding する
        "cat_idxs": list(categorical_index_with_modalities.keys()),
        "cat_dims": list(categorical_index_with_modalities.values()),
        # 埋め込み次元数はユニークな値の数から目安を決める
        "cat_emb_dim": [math.ceil(n ** 0.25) for n in categorical_index_with_modalities.values()],
        "optimizer_params": {
            "lr": 2e-3,
        },
    }

    # Unsupervised pre-training
    unsupervised_model = TabNetPretrainer(**tabnet_params)
    unsupervised_model.fit(
        X_train=train_x,
        eval_set=[
            eval_x,
        ],
        max_epochs=1_000,
        # 精度の改善が見られないときに Early stopping をかけるエポック数
        patience=100,
    )

    # Supervised fine-tuning
    clf = TabNetRegressor(**tabnet_params)
    clf.fit(
        X_train=train_x,
        y_train=train_y,
        eval_set=[
            (eval_x, eval_y),
        ],
        loss_fn=RMSLELoss(),
        eval_metric=[
            "rmsle",
        ],
        max_epochs=1_000,
        patience=100,
        from_unsupervised=unsupervised_model,
    )

    test_y_pred = clf.predict(test_x)

    # ホールドアウトしたテストデータで汎化性能を求める
    test_loss = rmsle_metric(test_y_pred, test_y)
    print(f"Test Loss: {test_loss:.8f}")

    # 特徴量の重要度を可視化する
    fig, ax = plt.subplots(1, 1)
    feature_imprtance = pd.Series(data={k: v for k, v in zip(x.columns, clf.feature_importances_)})
    feature_imprtance.sort_values().plot(kind="barh", ax=ax)
    ax.set_xlabel("importance")
    ax.set_ylabel("features")
    fig.savefig("importance.png")


if __name__ == '__main__':
    main()

上記を実行してみよう。 GPU (CUDA) は認識されていれば自動で使われる。 マシンの計算リソースによっては、だいぶ時間がかかる。

$ python tabnet_diamonds.py

...(snip)...

epoch 811| loss: 0.11108 | val_0_rmsle: 0.01169 |  0:20:33s
epoch 812| loss: 0.11114 | val_0_rmsle: 0.01161 |  0:20:34s
epoch 813| loss: 0.11284 | val_0_rmsle: 0.01422 |  0:20:36s

Early stopping occurred at epoch 813 with best_epoch = 713 and best_val_0_rmsle = 0.01049
/Users/amedama/.virtualenvs/py310/lib/python3.10/site-packages/pytorch_tabnet/callbacks.py:172: UserWarning: Best weights from best epoch are automatically used!
  warnings.warn(wrn_msg)
Test Loss: 0.10144833

テストデータに対して 0.10144833 という精度が得られた。

また、次のような特徴量の重要度を可視化したグラフが得られる。

TabNet の特徴量重要度

参考: LightGBM の場合のスコア

参考までに、勾配ブースティング決定木の代表として LightGBM を使った場合についても記載する。 あらかじめ LightGBM のパッケージをインストールする。

$ pip install lightgbm

サンプルコードは次のとおり。 LightGBM には学習用の損失関数として RMSLE が用意されていない。 そこで、代わりに目的変数を numpy.log1p() で変換した上で RMSE で学習している。

import numpy as np
import seaborn as sns
from sklearn.compose import ColumnTransformer
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import OrdinalEncoder
import lightgbm as lgb


def rmse_metric(y_pred, y_true):
    return mean_squared_error(y_pred, y_true, squared=False)


def main():
    df = sns.load_dataset('diamonds')
    x, y = df.drop(['price'], axis=1), df['price'].values

    y = np.log1p(y)

    categorical_cut_pipeline = Pipeline(steps=[
        ("cut-ordinal", OrdinalEncoder())
    ])
    categorical_color_pipeline = Pipeline(steps=[
        ("color-ordinal", OrdinalEncoder())
    ])
    categorical_clarity_pipeline = Pipeline(steps=[
        ("clarity-ordinal", OrdinalEncoder())
    ])
    numerical_pipeline = Pipeline(steps=[
        ("numerical", StandardScaler()),
    ])
    pipeline = ColumnTransformer(transformers=[
        ("categorical-cut-pipeline", categorical_cut_pipeline, ["cut"]),
        ("categorical-color-pipeline", categorical_color_pipeline, ["color"]),
        ("categorical-clarity-pipeline", categorical_clarity_pipeline, ["clarity"]),
        ("numerical-pipeline", numerical_pipeline, ["carat", "depth", "table", "x", "y", "z"])
    ])
    transformed_x = pipeline.fit_transform(x)

    train_x, test_x, train_y, test_y = train_test_split(
        transformed_x,
        y.reshape(-1, 1),
        test_size=0.33,
        random_state=42,
        shuffle=True,
    )

    train_x, eval_x, train_y, eval_y = train_test_split(
        train_x,
        train_y,
        test_size=0.2,
        random_state=42,
        shuffle=True,
    )

    lgb_train = lgb.Dataset(train_x, train_y)
    lgb_eval = lgb.Dataset(eval_x, eval_y, reference=lgb_train)

    lgb_params = {
        'objective': 'regression',
        'metric': "rmse",
        'verbose': -1,
        'seed': 42,
        'deterministic': True,
    }
    callbacks = [
        lgb.log_evaluation(period=10),
        lgb.early_stopping(
            stopping_rounds=100,
            first_metric_only=True,
        ),
    ]

    booster = lgb.train(
        params=lgb_params,
        train_set=lgb_train,
        valid_sets=[lgb_train, lgb_eval],
        num_boost_round=1_000,
        callbacks=callbacks,
    )

    test_y_pred = booster.predict(test_x,
                                  num_iteration=booster.best_iteration)

    test_loss = rmse_metric(test_y_pred, test_y)
    print(f"Test Loss: {test_loss:.8f}")


if __name__ == '__main__':
    main()

上記の実行結果は次のとおり。

$ python lgbm_diamonds.py

... (snip) ...

[700]  training's rmse: 0.0670463  valid_1's rmse: 0.0860752
[710]  training's rmse: 0.0668764  valid_1's rmse: 0.0860719
[720]  training's rmse: 0.0667363  valid_1's rmse: 0.0860665
Early stopping, best iteration is:
[623]  training's rmse: 0.0684748  valid_1's rmse: 0.0860188
Evaluated only: rmse
Test Loss: 0.08698780

テストデータに対して 0.08698780 という精度が得られた。

今回の条件においては LightGBM に分があるようだが、これはもちろんデータの性質などにも依存する。 また、性能を向上させる試みとして、勾配ブースティング決定木とニューラルネットワークをアンサンブルするのが常套手段である点も見逃すことはできない。

参考

github.com

arxiv.org

Linuxで動かしながら学ぶTCP/IPネットワーク入門

Linuxで動かしながら学ぶTCP/IPネットワーク入門

Amazon

  1. データの量がさほど多くないため、ニューラルネットワークが得意とするデータの性質ではないかもしれない↩