前処理ライブラリ、NVTabularを触ってみた - レバレジーズデータAIブログ

はじめに

こんにちは！レバレジーズデータ戦略室、データサイエンティストのJacobです。

今回は、大規模データ処理に特化したPythonライブラリ「NVTabular」について書きたいと思います。

どの機械学習モデルにも前処理が必要です。データセットを訓練・検証・テストセットに分割する前の前処理もあれば、分割後の前処理もあります。例えば、データ型変換や文字列の正規化は分割する前でも問題ないです。一方、連続変数の正規化・標準化やターゲットエンコーディングといった前処理は、訓練データから統計量やパラメータを学習し、それを検証・テストデータに適用します。これによりデータリーケージを防ぎ、本番環境での性能を正確に評価できます。

分割後の前処理は簡単なモデルであれば、データサイエンティストに長年愛用されているscikit-learnのPipelineでも問題ないですが、特徴量やデータ量が多いとパフォーマンスやメンテナンス性が課題になるかもしれません。今回はそれに対応するNVTabularというライブラリについて書きたいと思います。

NVTabular

NVTabularはNVIDIAがリリースした大規模データ処理ライブラリで、特にレコメンドシステムやテラバイト級の大容量データセットの前処理に特化しています。GPUとクラスターでの高速並列処理も可能です。

また、PyTorchやTensorFlowなどの深層学習フレームワークとの統合も容易で、前処理からモデル学習まで一貫したパイプラインを構築することが可能です。

コードサンプル

以下のようにWorkflowを定義します。scikit-learnのPipelineと同様の概念です。

次のような前処理をします。

カテゴリエンコーディング (Categorify): カテゴリ変数を整数値に変換します。
標準化 (Normalize): 連続変数の分布を平均0、分散1に変換します。
欠損値補完 (FillMedian): 欠損値を列の中央値で埋めます。
ターゲットエンコーディング (TargetEncoding): カテゴリ変数をターゲット変数の平均値を用いて数値に変換します。

import nvtabular as nvt

# カテゴリ変数の前処理
# 出現頻度が20回未満のカテゴリは「ヌルカテゴリ」にまとめる
CAT_COLS = ["hoge_col", "fuga_col"] # カテゴリカラム名
cat_feats = CAT_COLS >> nvt.ops.Categorify(freq_threshold=20)

# 連続変数の前処理
CONT_COLS = ["piyo_col"] # 連続変数カラム名
cont_feats = CONT_COLS >> nvt.ops.Normalize() >> nvt.ops.FillMedian()

# ターゲットエンコーディング
# 過学習を防ぐため、クロスバリデーションとスムージングを使う
TARGET_COLS = ["label"]  # ターゲットカラム名
target_encoded_feats = ["hoge_col"] >> nvt.ops.TargetEncoding(TARGET_COLS, kfold=5, p_smooth=20)

output = cat_feats + cont_feats + target_encoded_feats
workflow = nvt.Workflow(output)

前処理のDAGも可視化できます。

output.graph.render("dag", format="png")

その後は訓練セットに学習させて保存します。

train_dataset = nvt.Dataset("./train_dataset.parquet", engine="parquet")
workflow.fit(train_dataset)
workflow.save("./my_workflow")

推論する時はworkflowをロードし、前処理します。

workflow = nvt.Workflow.load("./my_workflow")
inference_dataset = nvt.Dataset("./inference_dataset.parquet", engine="parquet")
workflow.transform(inference_dataset).to_parquet('./processed_inference_dataset.parquet')

PyTorchで使いたい時は以下のデータローダが使えます。

from nvtabular.loader.torch import TorchAsyncItr

inference_dataset = nvt.Dataset("./processed_inference_dataset.parquet", engine="parquet")
data_loader = TorchAsyncItr(inference_dataset)
first_batch = next(iter(data_loader))