【AI Shift Advent Calendar 2021】pytorch用DataFrameライブラリTorchArrowを試してみる

こんにちは
AIチームの戸田です

本記事はAI Shift Advent Calendar 2021の8日目の記事です。

今回は先日プロトタイプが公開されました、pytorch用のDataFrameライブラリTorchArrowを試してみたいと思います

python用のDataFrameライブラリといいますと、pandasが有名ですね。最近だとKaggleなどのコンペティションでGPU対応のcudfもよく使われている印象です。

本記事では、TorchArrowの簡単な使い方を試したあと、pandasとcudfと速度比較を行います

TorchArrow

TorchArrowは、meta(旧facebook)が開発している深層学習のデータ前処理のためのtorch.TensorライクなDataFrameライブラリです。

現段階ではまだプロトタイプで、2022年頭にβ版をリリース予定だそうですが、プロトタイプのリポジトリのREADMEでは以下を実装予定だと言っています

  • ストリーミング処理をしやすくする
  • torch.DataLoaderなどのpytorchのAPIと簡単に連携できる
  • Arrowを使ったメモリの効率化
  • libcudfを使ったGPU対応
  • Veloxを使ったCPU処理の高速化
  • C++による高速なベクトル処理

GPUバックエンドにlibcudfを使っているので、ぱっと見cudfと変わらなそうですが、ストリーミング処理や、CPUでの高速化など、独自の強みがありそうです

使ってみる

サンプルデータとして、以下のような国数英のテストの点数を準備します

国語数学英語
佐藤609070
田中501020
中村403080
山田604040

このデータをTorchArrowのDataFrameにします

import torcharrow as ta

df = ta.DataFrame(
    {
        'Japanese': [60, 40, 50, 60],
        'Math': [90, 10, 30, 40],
        'English': [70, 20, 80, 40],
    }
)
df
'''
  index    Japanese    Math    English
-------  ----------  ------  ---------
      0          60      90         70
      1          40      10         20
      2          50      30         80
      3          60      40         40
dtype: Struct([Field('Japanese', int64), Field('Math', int64), Field('English', int64)]), count: 4, null_count: 0
'''

pandasと非常によく似ていますね。気をつけなければならないこととしては

  • 国語、のようなマルチバイト文字がつかえない
  • indexの指定ができない
  • .shapeで行列数を出せない

などがあり、このあたりは少し不便かな、と思います。β版だと解消されるかもしれないので、今後に期待です。

またta.DataFrameの引数としてdevice=”gpu”を指定するとGPUをしようすることができるようになるようですが、現在はまだ実装されていないようです

map関数やfilter関数もpandasのように使えるみたいです。詳しくはtutorialが公開されていますので、そちらをご参照ください

速度比較

速度比較として、上記で作ったDataFrameにおいて、数英の平均点を求める時間を計測します。データ数が少ないので、データを1000倍した際の時間で比較したいと思います

以下に速度計測のコードを示します

import pandas as pd

import torcharrow as ta
import torcharrow.dtypes as dt

import time
from contextlib import contextmanager

# 実行時間を計測する
@contextmanager
def timer(name):
    t0 = time.time()
    yield
    print(f'[{name}] done in {time.time() - t0:.5f} s')

# 英数の平均点を計算する
def avg_math_eng_1(a, b):
    return (a + b)/2

def avg_math_eng_2(row):
    return (row['Math'] + row['English'])/2

# データ作成
d = {
    'Japanese': [60, 40, 50, 60]*1000,
    'Math': [90, 10, 30, 40]*1000,
    'English': [70, 20, 80, 40]*1000,
}
ta_df = ta.DataFrame(d)
pd_df = pd.DataFrame(d)
cu_df = cudf.from_pandas(pd_df)


with timer('torcharrow'):
    ta_df.map(avg_math_eng_1, columns=['Math', 'English'], dtype=dt.float32)

with timer('pandas'):
    pd_df.apply(avg_math_eng_2, axis=1)

with timer('cudf'):
    cu_df.apply(avg_math_eng_2, axis=1)
TorchArrowpandascudf
実行時間(s)0.067780.041220.00225

まだプロトタイプだから、ということからかもしれないのですが、TorchArrowが一番遅い結果になってしまいました。

TorchArrowはストリーミング処理に強い、という事がREADMEに記載されていたので、ストリーミング処理ライクに、データを1000倍にするのではなく、同じデータを1000回流すように以下のようにコードを変更して、再度時刻を計測してみました。

d = {
    'Japanese': [60, 40, 50, 60],
    'Math': [90, 10, 30, 40],
    'English': [70, 20, 80, 40],
}
ta_df = ta.DataFrame(d)
pd_df = pd.DataFrame(d)
cu_df = cudf.from_pandas(pd_df)

with timer('torcharrow'):
    for _ in range(1000):
        ta_df.map(avg_math_eng_1, columns=['Math', 'English'], dtype=dt.float32)

with timer('pandas'):
    for _ in range(1000):
        pd_df.apply(avg_math_eng_2, axis=1)

with timer('cudf'):
    for _ in range(1000):
        cu_df.apply(avg_math_eng_2, axis=1)
TorchArrowpandascudf
実行時間(s)0.167580.463351.23027

今度はTorchArrowが最も早くなりました。やはりストリーミングに強いような実装になっているのだと思います。

cudfは大量のデータを一気に捌くのは得意ですが、小さいデータを何回も処理する際はGPUへの転送時間がネックになるようです。やりたいことに適した実装やライブラリを選ぶことが大切ですね。

終わりに

本記事では今回はpytorch用のDataFrameライブラリTorchArrowのプロトタイプを試して、pandasとcudfと現時点での実行時間比較を行いました。

よく使われるGroupByでの計算などでも比較したかったのですが、TorchArrowの方にまだ実装されていないようだったので、来年β版がリリースされたらまた試してみたいと思います

最後までお読みいただきありがとうございました!

明日はAmazonの反事実的な記述に関するデータセット、AMCDについての記事が公開される予定です