これだけ押さえればいい！dask_cudfの実装方法まとめ

はじめに

こんにちは、
株式会社オークンでデータサイエンティストをやっています、nobuです。

pandasでは計算が遅すぎるので、GPU計算ライブラリのcudfを使っていました。
cudfはpandasと同じような感覚で使えて手軽にGPU計算ができるという非常に優れたライブラリなのですが、少々問題があります。

cudfでは、GPU1枚に乗るまでのデータ容量でしか計算できないのです。

そこで、dask_cudfを使ったマルチGPUでの計算を導入したところ、大容量データの計算が可能になりました！

dask_cudfは実装に少し癖があり、導入にかなり苦労したのですが、
今回はこれだけ押さえればよい、dask_cudfの実装方法をまとめましたので、ご紹介できればと思います！

かなりざっくり、dask_cudfとは

dask_cudfとは、データを複数GPUに分散させながら、計算できる技術です。
dask_cudfを使うことで、1枚のGPUに乗り切らない大容量のデータでも、計算させることが可能になります。

さっそく、実装

まずは、事前準備として、以下のコードを書いておきます。これにより、利用可能なすべてのGPUを利用して計算できるようになります。

from dask_cuda import LocalCUDACluster

cluster = LocalCUDACluster()
client = Client(cluster)

データの読み込み

csvからも、parquetからも、データの読み込みが可能です。
できれば、parquetの方が読み込みが速いのでおすすめします。

import dask_cudf

ddf = dask_cudf.read_csv(filepath)
ddf = dask_cudf.read_parquet(filepath)

cudfからdask_cudfに変換することも可能です。

cudf = cudf.read_csv(filepath)
ddf = ddf.from_cudf(cudf, npartitions=n)

※ cudfからの変換だと、複数GPUにデータが分散されず、1つのGPUを集中して使ってしまうので、大容量のdfを変換する場合は注意が必要です。

daskの仕様

上記でdask_cudfへの変換を行いましたが、
実は、この時点では計算スケジュールが設定されただけで、読み込みは完了していません。

実際に読み込みを完了させるには、persistを使います。

ddf = ddf.persist()

persistを実行して初めて、各GPUにデータが分散されます。
ターミナル上で、watch nvidia-smiを実行して、GPUの使用具合を確認します。

今回はGPUが4枚搭載されたインスタンスを使ったのですが、
たしかに4枚のGPUにそれぞれデータが分散されていることが確認できます。

これは各種計算処理でも同様で、実際に計算させるにはpersistを使う必要があります。

persistは複数GPU上で演算させる処理ですが、
最終的な計算結果を取得する場合は、computeを実行する必要があります。

cudf = ddf.compute() # dask_cudfからcudfに変換されます
df = ddf.compute().to_pandas() # pandasに変換する場合

max_id = ddf['id'].max().compute() # 数値の計算結果を取得する場合も同様 

map_partitionについて

基本的にはpandasやcudfと同じように計算できますが、dask_cudfではサポートされていない計算が結構あります。
その場合は、map_partitionsを使うことで計算可能です。

import numpy as np

def myadd(df, n)
    return np.log(df['a']) * n
    
ddf["log_n"] = ddf.map_partitions(10) # n=10として渡される

RAPIDSコミュニティ

dask_cudfの導入には色々と躓いたのですが、なんとNVIDIA RAPIDSはSlackコミュニティを提供しているようです。
(RAPIDSのページの目立たないところにありました)

このSlackでは、RAPIDSの動作に関する質問を受け付けている？ようで、私も何度か質問に答えて頂きました!
RAPIDSを使われている方は是非活用されることをお勧めいたします。

最後に

無事、dask_cudfを導入することができ、メモリ不足に悩まされることがなく、大容量データでの計算が可能になりました！
大容量データを高速に計算したい場合はぜひご検討ください！
最後までご覧いただき、ありがとうございました！