Zenn
💽

Pytorch+Polarsで高速で動作するDatasetを作る

2023/12/01に公開
Python
Rust
PyTorch
Polars
tech

この記事はPolars Adbent Calendar 2023 2日目の記事です!!

2024/07/28追記

詳細な検証と原因を書いてくれた人がいたので、こちらを参照してください。
https://zenn.dev/uchiiii/articles/f58519345987ca

要約

Polarsにおいて行指定したイテレーション操作は非推奨なのだが、それでも圧倒的にPolarsが速いので、全人類はさっさとPolarsを使うこと。Pytorchのデータセットに使うと、pandasどころかnumpy製のものより高速で動作した。

初めに

Polars、みなさん使っていますか？筆者はPandasから乗り換えて四苦八苦して、ドキュメントとchatGPTにおんぶだっこでヒィヒィ言いながら使っています。

現状だと例えばraw dataの前処理や、EDAで使っています。ただ普段のML環境はPytorch-lightningとHydraで実装していることが多く、できればtorch.utils.data.Datasetやtorch.utils.data.DataLoaderクラス内でもPolarsで書き込みたいという事情があります。

大抵のデータ分割パッケージではインデックスを参照した分割を行うのですが、'Indexes are not needed!'と主張する偉大なる我らPolarsにはindex列などという前時代的な概念は存在しません（大本営発表）。迂回策としてはdatasetに投入する前にnumpyに変換すればよいというものがありますが、例えばマルチモーダルなモデルを作るために、テーブルデータと画像のパスが混在しているものをpl.DataFrameで一括で処理できた方がコード的にも理解しやすいと思っています。

ここまで考えたとき「インデックスをシャッフル分割したり、１行ずつ入出力する処理に関してはPolarsは実は苦手なんちゃうか？pandasにもまだ得意な領域は残されているんちゃうか？」という考えが浮かびました。

実験内容

今回はpytorch.Datasetをそれぞれのパッケージ(pandas/Polars/numpy)で提供されているDataFrameで実装した際の動作時間を計測します。

実験諸元

項目 内容
実行環境 MacOS(Ventura 13.5)+VSCode Notebook
Python3.10.8
torch==2.1.1
pandas==2.1.3
Polars==0.19.18
データサイズ 10000x3(feature2列, target1列)
KFold 3
バッチサイズ 4
エポック数 10

コード全体

import Polars as pl
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader, Subset
from sklearn.model_selection import KFold
import numpy as np
# ダミーデータ
rng = np.random.default_rng()

data = {
    "feature1": rng.random(10000),
    "feature2": rng.random(10000),
    "target": rng.random(10000),
}
pldf = pl.DataFrame(data)
pddf = pd.DataFrame(data)
nparray = np.array(
    [data["feature1"], data["feature2"], data["target"]]
).T


# Pandas Dataset
class PandasCustomDataset(Dataset):
    def __init__(self, dataframe):
        self.dataframe = dataframe

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, idx):
        features = self.dataframe.iloc[idx, :-1]
        target = self.dataframe.iloc[idx, -1]
        return torch.tensor(features, dtype=torch.float32), torch.tensor(
            target, dtype=torch.float32
        )


# Polars Dataset
class PolarsCustomDataset(Dataset):
    def __init__(self, dataframe):
        self.dataframe = dataframe

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, idx):
        features = np.array(self.dataframe.row(idx)[:-1])
        target = np.array(self.dataframe.row(idx)[-1])
        return torch.tensor(features, dtype=torch.float32), torch.tensor(
            target, dtype=torch.float32
        )


# Numpy Dataset(テーブルデータなら可能だが...)
class NumpyCustomDataset(Dataset):
    def __init__(self, ndarray):
        self.ndarray = ndarray

    def __len__(self):
        return len(self.ndarray)

    def __getitem__(self, idx):
        features = torch.Tensor(self.ndarray[idx, :-1])
        target = torch.Tensor([self.ndarray[idx, -1]])
        return torch.tensor(features, dtype=torch.float32), torch.tensor(
            target, dtype=torch.float32
        )


# データセットを作成
pldataset = PolarsCustomDataset(pldf)
pddataset = PandasCustomDataset(pddf)
npdataset = NumpyCustomDataset(nparray)

# Cross-validationのためのK-Foldを設定
dataset = pldataset
#dataset = pddataset
#dataset = npdataset
kf = KFold(n_splits=3)
for _fold, (train_index, valid_index) in enumerate(kf.split(range(len(dataset)))):
    train_dataset = Subset(dataset, train_index)
    valid_dataset = Subset(dataset, valid_index)

    batch_size = 4
    train_dataloader = DataLoader(train_dataset, batch_size, shuffle=True)
    valid_dataloader = DataLoader(valid_dataset, batch_size, shuffle=False)

    num_epochs = 10
    for i in range(num_epochs):
        for feats, target in train_dataloader:
            pass

実験結果

データセットの種類 タイム (秒)
pandas 13.3
Polars 1.6
numpy 2.7

OS環境によって多少はタイムの変動こそありますが、基本的にPolarsの方が爆速とみて間違いないと思います。意外だったのが最初からnumpy.ndarrayに変換したデータセットよりも速い点。コア部分がRustであることが関係してるんでしょうか。

References

この記事を書くために活用したサイトのURLはスクラップ記事一覧に記載しましたので、そちらを参考にしてください。

Discussion