torch.tensor への変換における Numpy と Polars の速度比較

これは、Polars Advent Calendar 2023 の 5日目です。
https://qiita.com/advent-calendar/2023/polars

はじめに

先日、Nista san によるブログが同じく Polars Advent Calendar 2023 の 2日目に発表されました。
https://zenn.dev/wotb_pythonista/articles/c5453b6e3d4625

上ブログは、着眼点が非常に興味深いです。確かに機械学習のモデルでは、column 単位でのアクセスではなく row 単位でのアクセスの頻度が多く、それが高速に動作することは非常に重要です。
row 単位の acesss が苦手な Polars が Pandas や Numpy と比較されていました。その結果は驚くことに、なんと Polars が一番。Numpy よりも高速に row 単位での random access ができたことになります。

これを読んだ時、驚きと同時に、いくら何でもそれは Polars 速すぎないかと少し気になってきました。というのも、Numpy は C言語で書かれており、その速度はもちろん、メモリ配列も row 単位でアクセスしやすいようになっているはず。加えて、何と言っても歴史があります。row 単位でのアクセスが苦手なぽっとでの Polars にあっさり負けてしまうのか、いやそんなはずないと思いました。（僕は Numpy の回し者ではないことに注意してください）

Polars が row 単位のアクセスが苦手な理由は、データを内部で Arrow format で保持しているからです。Arrow format に興味ある方は以下が公式のページです。
https://arrow.apache.org/docs/format/Columnar.html

そこで、このブログでは、上ブログで行われていたことをさらに詳細に検証してみようと思います。

結論

• Numpy と Polars を、純粋な row 単位の random access で比較すると、Numpy の方がかなり高速。
• しかし、numpy.array or Tuple -> torch.tensor の overhead が非常に大きいので、Numpy と Polars の row access の差は誤差になり、結果として Polars の方が高速に処理できるように見えていた。
• data サイズが大きくなる、つまり、row 数や column 数が大きくなると、どうなるかは未検証。

前提

以下の実験は次の環境で行なっています。

$ python --version
Python 3.9.16
$ pip list
polars                            0.19.19
pandas                            1.5.0
numpy                             1.23.3
torch                             1.12.1+cu116
...

特に断らない場合に、速度測定は、time(1) コマンドの結果（real）を用いており、計測は1回のみとします。（本来は複数回計測して平均を取ったりするべきですが、サボってます、すみません）

使用する Dummy data の大きさは全ての実験で統一してます。

data = {
    "feature1": rng.random(10000),
    "feature2": rng.random(10000),
    "target": rng.random(10000),
}

知識不足で誤ったことを書いてるかもしれませんが、その場合はご指摘いただけると嬉しいです。

Nista san のブログのコードを動かす

自分の環境で全く同じコードを動かして測定してみました。確かに、Polars が一番速度が短いです。やはり、驚きです。

コード

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

データセットの種類	タイム (sec)
polars	7.076
pandas	32.539
numpy	9.181

純粋な random access の速度比較

以下のコードで torch などを用いずに純粋な random access の速度を比較しました（time(1) コマンドは使ってません）

コード

import polars as pl
import pandas as pd
import numpy as np

from functools import wraps
import time

def timeit(func):
    @wraps(func)
    def timeit_wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        total_time = end_time - start_time
        print(f'Function {func.__name__}{args} {kwargs} Took {total_time:.4f} seconds')
        return result
    return timeit_wrapper

rng = np.random.default_rng()

data = {
    "x1": rng.random(10000),
    "x2": rng.random(10000),
    "y": rng.random(10000),
}
pldf = pl.DataFrame(data)
pddf = pd.DataFrame(data)
nparray = np.array([data['x1'], data['x2'], data['y']]).T
num_cols = 3

random_indices = rng.integers(0, 10000, 1000000)

@timeit
def _polars():
    for index in random_indices:
        for col in range(num_cols):
            _ = pldf.row(index)[:-1]

@timeit
def _pandas():
    for index in random_indices:
        for col in range(num_cols):
            _ = pddf.iloc[index, :-1]

@timeit
def _numpy():
    for index in random_indices:
        for col in range(num_cols):
            _ = nparray[index, :-1]

if __name__=="__main__":
    _polars()
    _pandas()
    _numpy()

結果は以下。

library	タイム (sec)
polars	2.8568
pandas	190.9075
numpy	0.8303

Numpy の方が 3 倍以上高速ですね。やはり、Numpy と比較して、Polars は row 単位での random access が不得意なことがわかります。（と言っても、Pandas よりは全然はやいです）
早速先ほどの実験の結果と矛盾する結果が得られました。面白くなってきましたね。

Nista san のコードで torch.Tensor ではなく np.array を返すようにするとどうなる？

Dataset の __getitem__ の返り値を torch.Tensor ではなく np.array にするとどうなるのか。

コード

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),
}


# 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 features, target


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

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

    def __getitem__(self, idx):
        features = self.ndarray[idx, :-1]
        target = np.array(self.ndarray[idx, -1])
        return features, target


class DummyCustomDataset(Dataset):
    def __init__(self):
        self.dummy_features = np.ones(3)
        self.dummy_target = np.array(1.0)

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

    def __getitem__(self, idx):
        return self.dummy_features, self.dummy_target

def _polars():
    pldf = pl.DataFrame(data)
    return PolarsCustomDataset(pldf)


def _pandas():
    pddf = pd.DataFrame(data)
    return PandasCustomDataset(pddf)


def _numpy():
    nparray = np.array(
        [data["feature1"], data["feature2"], data["target"]]
    ).T
    return NumpyCustomDataset(nparray)


def _dummy():
    return DummyCustomDataset()

#dataset = _polars()
dataset = _pandas()
#dataset = _numpy()
#dataset = _dummy()

random_indices = rng.integers(0, 10000, 1_000_000)

for index in random_indices:
    _ = dataset[index]

結果は以下です。

データセットの種類	タイム (sec)
polars	4.625
pandas	140.207
numpy	2.761
dummy	2.146

ここで、dummy とは常に同じ np.array を返す Dataset のことを指します。（コード参照）
驚くことに、この dummy ですら 2.146 sec かかっています。dummy は random access の部分には関係なく必要な大まかな時間だと思って問題ないでしょう。なので、この時間をそれぞれから引いてやります。

データセットの種類	タイム (sec)
polars	2.479
pandas	138.061
numpy	0.615
dummy	0

Polars は Numpy の4倍弱かかってますね。純粋な random access の速度比較 の結果と整合してます。しかし、Nista san のブログのコードを動かした結果とは辻褄が合いません。
もしかすると、random access 以外のところで、時間を多く費やしてる部分がありそうです。

以降では、他の比べて圧倒的に速度が遅く特に比較する必要がないため、Pandas の結果は除きます。

ちょこっと修正 (ver1)

コードを読んでると少し気になる部分がありました。以下の NumpyCusomDataset の __getitem__ 関数です。
内部で、numpy.array を torch.Tensor に変換し、それをもう一度 torch.Tensor に変換して返却してます。Polars などでは、tensor.Tensor が作られるのは1度のみなのでそこの差分がありそうです。また、target は List にすることなくそのまま torch.Tensor に変えてやることで heap allocation を防げそうです。

@@ -60,22 +60,38 @@ class NumpyCustomDataset(Dataset):
         return len(self.ndarray)
 
     def __getitem__(self, idx):
-        features = torch.Tensor(self.ndarray[idx, :-1])
-        target = torch.Tensor([self.ndarray[idx, -1]])
+        features = self.ndarray[idx, :-1]
+        target = self.ndarray[idx, -1]
         return torch.tensor(features, dtype=torch.float32), torch.tensor(
             target, dtype=torch.float32
         )

Dummy の　Dataset も加えておきます。

+class DummyCustomDataset(Dataset):
+    def __init__(self):
+        self.dummy_features = np.ones(3)
+        self.dummy_target = np.array(1.0)
+
+    def __len__(self):
+        return 10000
+
+    def __getitem__(self, idx):
+        return torch.tensor(self.dummy_features, dtype=torch.float32), torch.tensor(
+            self.dummy_target, dtype=torch.float32
+        )

さて、上記の変更を加えた上でもう一度計測してみましょう。

データセットの種類	タイム (sec)
polars	6.982
numpy	6.216
dummy	5.914

お！Numpy の方が Polars より高速になった！

さらに、dummy との差分を取ると、

データセットの種類	タイム (sec)
polars	1.068
numpy	0.302
dummy	0

となります。Polars は Numpy の 3 倍以上 random access にかかっていそうですね。なるほど、torch.Tensor への変換がかなり時間を消費していたようですね。これであれば、先ほどの結果と辻褄が合います。
それにしても Dummy が 5.914 sec かかっていることから、np.array -> torch.Tensor への変換が非常に重いですね。（どうしてかわかる人教えてください!）

ちょこっと修正 (ver2)

先ほどの修正に加えて、さらに修正してみます。
また、PolarsCusomDataset に関しては、np.array に変換するより直接 torch.Tensor に変換したほうが無駄な overhead は減りそうです。

ちなみに、polars.DataFrame.row 関数は default で Tuple を返します。
https://pola-rs.github.io/polars/py-polars/html/reference/dataframe/api/polars.DataFrame.row.html

@@ -44,8 +44,8 @@ class PolarsCustomDataset(Dataset):
         return len(self.dataframe)
 
     def __getitem__(self, idx):
-        features = np.array(self.dataframe.row(idx)[:-1])
-        target = np.array(self.dataframe.row(idx)[-1])
+        features = self.dataframe.row(idx)[:-1]
+        target = self.dataframe.row(idx)[-1]
         return torch.tensor(features, dtype=torch.float32), torch.tensor(
             target, dtype=torch.float32
         )

これによって、PolarsCustomDataset に関しては np.array -> torch.tensor ではなく、Tuple -> torch.tensor の変換になります。なので、それに対応する Dummy も追加します。（dummytuple と呼びます）

+class DummyTupleCustomDataset(Dataset):
+    def __init__(self):
+        self.dummy_features = (1.0, 1.0, 1.0)
+        self.dummy_target = 1.0
+
+    def __len__(self):
+        return 10000
+
+    def __getitem__(self, idx):
+        return torch.tensor(self.dummy_features, dtype=torch.float32), torch.tensor(
+            self.dummy_target, dtype=torch.float32
+        )

その結果がこちら。

データセットの種類	タイム (sec)
polars	5.752
numpy	6.216
dummy	5.914
dummytuple	5.172

あれ、Polars の方が高速になった！さらに、dummy どうしを比較することで、np.array -> torch.tensor よりも tuple -> torch.tensor の方が高速なことがわかります！（理由ご存知の方いれば教えてください！）

しかし、それぞれ 比較対象の dummy との差分をとってやると

データセットの種類	タイム (sec)
polars	0.58
numpy	0.302

やはり、random access 部分の速度は　Numpy の方が上ですね！
一見 Polars の方が高速に見えていたのは、 Polars そのものの処理速度ではなく、torch.tersor に変換する時の処理の速度によるものだったと結論づけて良さそうです。

では、以下のように変更して再度測定してみます。結果は numpy (with tuple) とします。

@@ -60,22 +60,53 @@ class NumpyCustomDataset(Dataset):
         return len(self.ndarray)
 
     def __getitem__(self, idx):
-        features = torch.Tensor(self.ndarray[idx, :-1])
-        target = torch.Tensor([self.ndarray[idx, -1]])
+        features = tuple(self.ndarray[idx, :-1])
+        target = self.ndarray[idx, -1]
         return torch.tensor(features, dtype=torch.float32), torch.tensor(
             target, dtype=torch.float32
         )

データセットの種類	タイム (sec)
polars	5.752
numpy	6.216
numpy (with tuple)	5.893
dummy	5.914
dummytuple	5.172

なるほど、これだとまだ Polars と同じか遅いくらいですね。np.array が heap にありそれを copy してくるのに時間がかかっているのかもしれません。

ちょこっと修正 (ver3)

torch にこんな関数があることを同僚の方に教えていただきました。
https://pytorch.org/docs/stable/generated/torch.from_numpy.html

これは numpy の memory をコピーせずそのまま torch 側で使えるようにするもののようです。注意点として、resize などはできないと書かれてます。
これを使えば、numpy の速度向上が見込めそうです。試してみましょう。numpy (with torch.from_numpy) とします。

@@ -60,22 +60,53 @@ class NumpyCustomDataset(Dataset):
         return len(self.ndarray)
 
     def __getitem__(self, idx):
         features = self.ndarray[idx, :-1]
         target = self.ndarray[idx, -1]
-        return torch.tensor(features, dtype=torch.float32), torch.tensor(
+        return torch.from_numpy(features), torch.tensor(
             target, dtype=torch.float32
         )

データセットの種類	タイム (sec)
polars	5.752
numpy	6.216
numpy (with tuple)	5.893
numpy (with torch.from_numpy)	5.067
dummy	5.914
dummytuple	5.172

おお、かなり高速になりました。Polars よりもかなり優位ですね！
あれ、zero copy になったはずなのに、まだ 5 sec かかってます。全く関係ないところでかなりの時間使ってそうですね。

どんどん踏み込んでいきたいところですが、一旦ここら辺でまとめておきます。

結論（再掲）

• Numpy と Polars を、純粋な row 単位の random access で比較すると、Numpy の方がかなり高速。
• しかし、numpy.array or Tuple -> torch.tensor の overhead が非常に大きいので、Numpy と Polars の row access の差は誤差になり、結果として Polars の方が高速に処理できるように見えていた。
• data サイズが大きくなる、つまり、row 数や column 数が大きくなると、どうなるかは未検証。

ただし、あくまでこの検証は　Polars が苦手な row 単位での access の話です！！！！
column 単位でのアクセスなどは、驚くほどに高速なので、みなさんぜひ Polars を使ってみてその速度を体感してみてください！！！