こんにちは
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での高速化など、独自の強みがありそうです
使ってみる
サンプルデータとして、以下のような国数英のテストの点数を準備します
| 国語 | 数学 | 英語 | |
| 佐藤 | 60 | 90 | 70 |
| 田中 | 50 | 10 | 20 |
| 中村 | 40 | 30 | 80 |
| 山田 | 60 | 40 | 40 |
このデータを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)| TorchArrow | pandas | cudf | |
| 実行時間(s) | 0.06778 | 0.04122 | 0.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)| TorchArrow | pandas | cudf | |
| 実行時間(s) | 0.16758 | 0.46335 | 1.23027 |
今度はTorchArrowが最も早くなりました。やはりストリーミングに強いような実装になっているのだと思います。
cudfは大量のデータを一気に捌くのは得意ですが、小さいデータを何回も処理する際はGPUへの転送時間がネックになるようです。やりたいことに適した実装やライブラリを選ぶことが大切ですね。
終わりに
本記事では今回はpytorch用のDataFrameライブラリTorchArrowのプロトタイプを試して、pandasとcudfと現時点での実行時間比較を行いました。
よく使われるGroupByでの計算などでも比較したかったのですが、TorchArrowの方にまだ実装されていないようだったので、来年β版がリリースされたらまた試してみたいと思います
最後までお読みいただきありがとうございました！
明日はAmazonの反事実的な記述に関するデータセット、AMCDについての記事が公開される予定です
