🤖

PyTorch TabularとPyTorch Frameの比較

に公開

概要

この記事では、表形式データの深層学習モデルを扱うライブラリとして、PyTorch TabularとPyTorch Frameを比較します。両者は設計思想が異なるため、目的(素早いベースライン作成/拡張性の高い実装)に応じて選択します。

PyTorch Tabularとは?

PyTorch Tabularとは、表形式データを扱う深層学習モデルを扱うためのライブラリです。学習やsweep機能によって、比較的容易にモデルを構築できるのが特徴です。主な機能は以下のとおりです。

  • TabNetやDANETsといったモデルを共通のAPIを通じて利用可能
  • ハイパーパラメータ探索やモデルスイープが可能
  • モデルや損失関数、評価関数のカスタム設定が可能
  • Automatic Feature Interaction Network、FTTransformer、TabNet等が利用可能

PyTorch Tabularは、様々なモデルをすぐに試したいといった人に向いたライブラリになっています。

PyTorch Frameとは?

PyTorch Frameは、異なるデータタイプカラムが存在している表形式のデータでも容易にモデル設計を可能としたライブラリです。PyTorchのように学習および推論を行うのが特徴です。
主な機能は以下のとおりです。

  • Encoder・Decoder設計による多様なデータタイプへの対応
  • PyTorchのような部品設計が可能
  • FTTransformer、ExcelFormer、TabNet等が利用可能。PyTorch Tabularよりは扱えるモデルが少ない

PyTorch Frameは特にモデルの新規開拓や異なるデータタイプの扱いを詳細に設定できるのが特徴です。

比較

それぞれのライブラリは、どちらも深層学習モデルを利用し、表形式データを扱うためのライブラリです。ただ、コード設計が大きく異なるため、それらを紹介していきます。
今回はデータセットとして、Mercari Price Suggestion Challengeを利用します。

全体のコード設計

まず、それぞれのライブラリをインストールします。それぞれのライブラリは、pipでインストールできます。

!pip install pytorch-frame pytorch-tabular

また、以下のコードによって、データを取得

!apt-get install p7zip
!p7zip -d -f -k /kaggle/input/mercari-price-suggestion-challenge/train.tsv.7z

その後、

import pandas as pd
df_train = pd.read_csv("/kaggle/working/train.tsv", sep="\t")
train_dropped = df_train.dropna()
train_dropped = train_dropped.drop(columns=["name", "category_name", "brand_name"])
train_dropped = train_dropped.sample(frac=0.01, random_state=1).reset_index(drop=True)

によって、以下のような表が得られます。

train_id item_condition_id price shipping item_description
545540 1 50.0 0 LAST ONE. PRICE FIRM. New in sealed box Raptur...
400100 2 14.0 0 Inside of waistband has a cute print on!
412785 3 25.0 1 I only wore this once, because it was too smal...
121615 3 35.0 1 The unit has been tested and works. It has som...
59683 1 28.0 1 Victoria's Secret Very Sexy fuchsia pink longl...

今回は、処理に簡略化のためにデータサイズを縮小しています。

データを見ると、表の中にもテキストデータが混ざっています。このデータにおいて、目的変数がpriceになります。
このデータを利用して、PyTorch TabularとPyTorch Frameで実装で回帰予測してみます。なお、ベースとなるモデルはどちらもFTTransformerとします。

PyTorch Tabularのコード

PyTorch Tabularでテキストを扱うために、今回はSentence Transformerで埋め込みを得て処理していきます。埋め込みの取得に、DistilBERT(distilbert-base-uncased)を利用し、処理します。

## 必要なライブラリ
import torch
import numpy as np
from pytorch_tabular import TabularModel
from pytorch_tabular.config import DataConfig, OptimizerConfig, TrainerConfig
from pytorch_tabular.models import FTTransformerConfig
from sentence_transformers import SentenceTransformer
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

def encode_col(df, col, model, batch_size=128, normalize=False):
    """
    埋め込みの取得。SentenceTransformerを利用
    """
    texts = df[col].astype(str).tolist()
    with torch.no_grad():
        embs = model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=True,
            convert_to_numpy=True,
            normalize_embeddings=normalize,
        )
    return embs  # (N, dim)

def add_emb_cols(df, embs, prefix):
    """
    埋め込みの結合
    """
    embs = np.asarray(embs)
    dim = embs.shape[1]
    emb_df = pd.DataFrame(
        embs,
        index=df.index,
        columns=[f"{prefix}_{i}" for i in range(dim)],
    )
    
    return pd.concat([df, emb_df], axis=1)

model = SentenceTransformer("distilbert-base-uncased", device="cuda" if torch.cuda.is_available() else "cpu")

tabular_df = train_dropped.copy()
desc_emb  = encode_col(tabular_df, "item_description", model, batch_size=128)
tabular_df = tabular_df.drop(columns=["item_description"])

tabular_df = add_emb_cols(tabular_df, desc_emb,  "item_description_emb")
df_train, df_eval_test = train_test_split(tabular_df, test_size=0.2, random_state=1)
df_eval, df_test = train_test_split(df_eval_test, test_size=0.5, random_state=1)

上記のコードによって埋め込みを得て、データを作成した後に、以下のコードで学習と評価:

# データの設定。連続値と目的変数を設定する
data_config = DataConfig(
    target=["price"],
    continuous_cols=[
        col for col in df_train.columns if col not in ["price"]
    ],
)
# 学習設定。今回は設定していませんが EarlyStopping などの設定が可能
trainer_config = TrainerConfig(
    auto_lr_find=True,
    batch_size=64,
    max_epochs=20,
)
# 最適化関数の設定
optimizer_config = OptimizerConfig(optimizer="AdamW")

# FTTransformerの設定。モデル設計はデフォルト値を利用
model_config = FTTransformerConfig(
    task="regression",
    learning_rate=1e-3,
)
# Tabularの設定
tabular_model = TabularModel(
    data_config=data_config,
    model_config=model_config,
    optimizer_config=optimizer_config,
    trainer_config=trainer_config,
)
# 学習
tabular_model.fit(train=df_train, validation=df_eval)
# 推論
pred_df = tabular_model.predict(df_test)
# 評価
print(
    mean_squared_error(
        df_test["price"].values,
        pred_df["price_prediction"].values,
    )
)

コード自体はシンプルで、データ、モデルと学習の設定を行い、fitとpredictするだけで簡単に予測することができます。また、モデル自体も柔軟に変更が可能であり、簡単に実験がしやすそうです。学習中の出力は以下のとおりです。

Trainable params: 323 K                                                                                            
Non-trainable params: 0                                                                                            
Total params: 323 K                                                                                                
Total estimated model params size (MB): 1                                                                          
Modules in train mode: 122                                                                                         
Modules in eval mode: 0                                                                                            
Total FLOPs: 0                                                                                                     
Epoch 19/19 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 212/212 0:01:02 • 0:00:00 3.41it/s v_num: 2.000 train_loss: 1142.746 
                                                                                 valid_loss: 2826.071              
                                                                                 valid_mean_squared_error: 2826.071

予測の評価結果は以下のとおりです
MSE: 2759.0129

PyTorch Frameのコード

PyTorch FrameではPyTorch Tabularのように事前に埋め込みを得てから学習する方法および、end-to-endで実装する方法があります。それぞれには性能差があるのですが、変更は比較的容易です。
まず、必要なライブラリをimportします。

import torch.nn.functional as F

from torch_frame.data import Dataset, DataLoader
from torch_frame import stype
from torch_frame.utils import infer_df_stype
from torch_frame.utils.split import generate_random_split

from torch_frame.nn import EmbeddingEncoder, LinearEncoder, LinearEmbeddingEncoder
from torch_frame.nn.models.ft_transformer import FTTransformer
from torch_frame.config.text_embedder import TextEmbedderConfig
from torch_frame.config.text_tokenizer import TextTokenizerConfig

次に、事前に処理するパターンでは以下の通り。PyTorch FrameはデータにTensorFrameを利用します。

# 事前処理バージョン
class TextToEmbedding:
    def __init__(self, device):
        self.model = SentenceTransformer("distilbert-base-uncased", device=device)

    def __call__(self, sentences):
        emb = self.model.encode(sentences, convert_to_tensor=True)
        return emb.cpu() 

all_df = train_dropped.copy()
# 学習・検証・テストの分割を設定。テストは`include_test`で設定できる(デフォルトはTrue)
all_df["split"] = generate_random_split(len(all_df), seed=1, train_ratio=0.8, val_ratio=0.1)

col_to_stype = infer_df_stype(all_df) # 自動で各列のstypeを推定
col_to_stype["price"] = stype.numerical # 目的変数(連続値)
col_to_stype.pop("split") # Dataset作成時に含めない

# テキストデータの設定
text_cols = ["item_description"]
for col in text_cols:
    col_to_stype[col] = stype.text_embedded

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 埋め込みの設定
col_to_text_embedder_cfg = TextEmbedderConfig(
    text_embedder=TextToEmbedding(device),
    batch_size=64,
)

# データセットの定義
dataset = Dataset(
    all_df,
    col_to_stype=col_to_stype,
    target_col="price",
    split_col="split",
    col_to_text_embedder_cfg=col_to_text_embedder_cfg,
)
# materialize
dataset.materialize()

train_ds, val_ds, test_ds = dataset.split()
train_loader = DataLoader(train_ds.tensor_frame, batch_size=64, shuffle=True)
val_loader   = DataLoader(val_ds.tensor_frame, batch_size=128, shuffle=False)

# Encoderの設定
stype_encoder_dict = {
    stype.categorical: EmbeddingEncoder(),
    stype.numerical: LinearEncoder(),
    stype.embedding: LinearEmbeddingEncoder(), 
}

# モデル設定
model = FTTransformer(
    channels=64,
    out_channels=1,          # 回帰
    num_layers=4,
    col_stats=train_ds.col_stats,  
    col_names_dict=train_ds.tensor_frame.col_names_dict,
    stype_encoder_dict=stype_encoder_dict,
).to(device)

opt = torch.optim.AdamW(model.parameters(), lr=3e-4)

その後、学習と評価関数を定義:

def train_one_epoch(model, loader, opt):
    """
    1epoch学習する関数
    """
    model.train()
    total = 0.0
    for tf in loader:
        tf = tf.to(device)
        pred = model(tf).squeeze(-1)
        y = tf.y.float()
        loss = F.mse_loss(pred, y)
        opt.zero_grad()
        loss.backward()
        opt.step()
        total += loss.item() * len(y)
    return total / len(loader.dataset)

@torch.no_grad()
def eval_mse(model, loader):
    """
    評価関数
    """
    model.eval()
    total = 0.0
    for tf in loader:
        tf = tf.to(device)
        pred = model(tf).squeeze(-1)
        y = tf.y.float()
        loss = F.mse_loss(pred, y)
        total += loss.item() * len(y)
    return total / len(loader.dataset)

これらの関数によって、以下によって学習と評価:

for epoch in range(1, 21):
    train_loss = train_one_epoch(model, train_loader, opt)
    val_loss = eval_mse(model, val_loader)
    print(f"Epoch {epoch:02d}: Train MSE={train_loss:.4f}, Val MSE={val_loss:.4f}")

print("Evaluating on test set...")
test_loader = DataLoader(test_ds.tensor_frame, batch_size=128, shuffle=False)
test_loss = eval_mse(model, test_loader)

コード自体は、PyTorchを利用した学習に似ていて、推論→損失の計算→backwardの流れを自身で記述します。PyTorch Tabularほど簡単に記述できるわけではないのですが、カスタム性が高く、実験的な用途に向いています。学習中の出力は以下のとおりです:

Epoch 01: Train MSE=3057.1089, Val MSE=3227.2752
Epoch 02: Train MSE=2964.1568, Val MSE=3154.1341
Epoch 03: Train MSE=2890.4000, Val MSE=3079.0465
Epoch 04: Train MSE=2815.5521, Val MSE=3001.9942
Epoch 05: Train MSE=2740.6352, Val MSE=2925.9750
Epoch 06: Train MSE=2667.2354, Val MSE=2852.2251
Epoch 07: Train MSE=2596.5819, Val MSE=2781.7084
Epoch 08: Train MSE=2530.0745, Val MSE=2716.3673
Epoch 09: Train MSE=2468.9841, Val MSE=2656.2359
Epoch 10: Train MSE=2414.4366, Val MSE=2602.1641
Epoch 11: Train MSE=2367.0606, Val MSE=2555.8676
Epoch 12: Train MSE=2326.7244, Val MSE=2516.4416
Epoch 13: Train MSE=2293.3433, Val MSE=2484.5181
Epoch 14: Train MSE=2266.2281, Val MSE=2458.5733
Epoch 15: Train MSE=2245.1347, Val MSE=2438.2614
Epoch 16: Train MSE=2229.0530, Val MSE=2423.2945
Epoch 17: Train MSE=2217.3837, Val MSE=2411.7032
Epoch 18: Train MSE=2209.0269, Val MSE=2403.5408
Epoch 19: Train MSE=2203.2613, Val MSE=2397.7993
Epoch 20: Train MSE=2199.4328, Val MSE=2393.9080

評価結果は以下のとおりです。
MSE:1515.4272

また、end-to-endにしたい場合は、以下のクラスを定義:

from torch import nn
from transformers import AutoModel
from torch_frame.config import ModelConfig
from torch_frame.nn.encoder import LinearModelEncoder

class TextToEmbeddingTokenization:
    def __init__(self, name="tohoku-nlp/bert-base-japanese-v3"):
        self.tok = AutoTokenizer.from_pretrained(name)

    def __call__(self, sentences: list[str]):
        return self.tok(
            sentences,
            truncation=True,
            padding=True,
            return_tensors="pt",
        )

class TextToEmbeddingFinetune(nn.Module):
    def __init__(self, name="distilbert-base-uncased"):
        super().__init__()
        self.model = AutoModel.from_pretrained(name) # transformersを利用

    def forward(self, feat):
        input_ids = feat["input_ids"].to_dense(fill_value=0).squeeze(1)
        attn_mask = feat["attention_mask"].to_dense(fill_value=0).squeeze(1)

        out = self.model(input_ids=input_ids, attention_mask=attn_mask)
        cls = out.last_hidden_state[:, 0, :]          # [B, H]
        return cls.unsqueeze(1)                        # [B, 1, H]

して、col_to_stypeの設定およびcfgの設定を変更します。

...
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
text_cols = ["item_description"]
for col in text_cols:
    col_to_stype2[col] = stype.text_tokenized # 変更箇所1 text_tokenizedにする

col_to_text_tokenizer_cfg = TextTokenizerConfig( # 変更箇所2
    text_tokenizer=TextToEmbeddingTokenization(),
    batch_size=10_000,
)
col_to_text_tokenizer_cfg = {col: tok_cfg for col in text_cols}

dataset2 = Dataset(
    all_df2,
    col_to_stype=col_to_stype2,
    target_col="price",
    split_col="split",
    col_to_text_tokenizer_cfg=col_to_text_tokenizer_cfg, # 変更箇所3
)
...
# 埋め込み用のモデル微調整の設定 (変更箇所4)
text_model = TextToEmbeddingFinetune("distilbert-base-uncased").to(device)
col_to_model_cfg = {
    col: ModelConfig(model=text_model, out_channels=768)
    for col in text_cols
}
stype_encoder_dict2 = {
    stype.categorical: EmbeddingEncoder(),
    stype.numerical: LinearEncoder(),
    stype.text_tokenized: LinearModelEncoder(col_to_model_cfg=col_to_model_cfg),
}

model2 = FTTransformer(
    channels=64,
    out_channels=1,          # 回帰
    num_layers=4,
    col_stats=train_ds2.col_stats,  
    col_names_dict=train_ds2.tensor_frame.col_names_dict,
    stype_encoder_dict=stype_encoder_dict2,
).to(device)
...

あとは同様のコードを実行することで学習することができます。学習の結果、評価値は以下の通り:
MSE:1501.5943

全体として、PyTorch Frameを利用した方がMSEが小さくなりました。これは、前処理(テキスト特徴量の作り方)や学習設定が異なるため,生じたかと思います。PyTorch Frameはこうした詳細設計が可能であるので、カスタマイズ性はよさそうです。

まとめ

今回は、PyTorch TabularとPyTorch Frameを比較してみました。どちらも表形式データを扱う深層学習モデルを利用するために便利なライブラリですが、PyTorch Tabularは実験が容易な点、PyTorch Frameはカスタマイズ性が高いです。それぞれには違いが生じているので、目的に応じて使い分けてみてください。

Discussion