pandasは遅い！大量データの高速化、最適化のノウハウを公開！

riiidコンペ(kaggle)お疲れさまでした！

今回は、銅メダルの取得で終わりましたが、学びが多かったコンペとも言えます。

pandasは、条件を指定してデータを取り出したりするのに便利なライブラリですが、その中でも特に苦戦した、大量データの扱い、pandasの処理の高速化について記事化したいと思います！

riiidのデータはどんなデータだったの？

riiidコンペは、ざっくり言うとTOEIC学習アプリで使われているAIの精度を上げるものです。

コンペの詳細は、以下のノートブックにまとめたので参照してください！

【日本語】[Japanese] Riiid コンペに取り組む前の準備・随時更新

riiidコンペのデータ(CSVデータ)は、1億レコード(行)もあり、単純にメモリに乗せると4Gほどの大きなデータになります。更に、特徴量を追加したりすると8G,16Gになってしまいすぐにメモリ不足になってしまいます。

メモリを効果的に使うスキルが求められるコンペといえました。

メモリ容量の最適化

pandasにread_csvだけでデータを読み込むだけでは、無駄な領域まで使ってしまい、すぐにメモリ不足に陥ります。

ここでは、pandasにおいてできるかぎり少ないメモリ容量で処理できる方法を書きます。

大前提として、使用しないメモリ(変数)が出てきたら、その都度、delで解放していきます。

データに対して型を宣言する

単純にpandasにread_csvしただけでは、最もデータ領域を使う型が選択されてしまいます。たとえば、int系だとint64が使われてしまいます。

数値の範囲が、127までの場合、int16で足りるのに、int64を使っていたら無駄な領域使うことになります。

こんなときは、型を定義して、「astype」関数を組み合わせるなどして、メモリ容量を減らすことを検討します。

コードの例は以下の通りです。

data_types_dict = {
    'timestamp': 'int64',
    'user_id': 'int32', 
    'content_id': 'int16', 
    'content_type_id':'int8', 
    'task_container_id': 'int16',
    'user_answer': 'int8',
    'answered_correctly': 'int8', 
    'prior_question_elapsed_time': 'float32', 
    'prior_question_had_explanation': 'bool'
}

#データは以下からダウンロードしてください
#https://www.kaggle.com/c/riiid-test-answer-prediction/data
train_df = pd.read_csv('./riiid-test-answer-prediction/train.csv',nrows = 1000000)

#型を宣言しないと、無駄にメモリを使う(例：本来int8で足りるのに、int64で確保される)のでしっかり型宣言しておく
train_df = train_df.astype(data_types_dict)

関数を作りメモリ容量を最適化する処理をする

以下のような関数を使ってメモリ容量を最適化します。

'''Function to reduce the DF size'''
# source: https://www.kaggle.com/kernels/scriptcontent/3684066/download

def reduce_mem_usage(df):
    """ iterate through all the columns of a dataframe and modify the data type
        to reduce memory usage.        
    """
    start_mem = df.memory_usage().sum() / 1024**2
    print('Memory usage of dataframe is {:.2f} MB'.format(start_mem))
    
    for col in df.columns:
        col_type = df[col].dtype
        
        if col_type != object:
            c_min = df[col].min()
            c_max = df[col].max()
            if str(col_type)[:3] == 'int':
                if c_min > np.iinfo(np.int8).min and c_max < np.iinfo(np.int8).max:
                    df[col] = df[col].astype(np.int8)
                elif c_min > np.iinfo(np.int16).min and c_max < np.iinfo(np.int16).max:
                    df[col] = df[col].astype(np.int16)
                elif c_min > np.iinfo(np.int32).min and c_max < np.iinfo(np.int32).max:
                    df[col] = df[col].astype(np.int32)
                elif c_min > np.iinfo(np.int64).min and c_max < np.iinfo(np.int64).max:
                    df[col] = df[col].astype(np.int64)  
            else:
                if c_min > np.finfo(np.float16).min and c_max < np.finfo(np.float16).max:
                    df[col] = df[col].astype(np.float16)
                elif c_min > np.finfo(np.float32).min and c_max < np.finfo(np.float32).max:
                    df[col] = df[col].astype(np.float32)
                else:
                    df[col] = df[col].astype(np.float64)
        else:
            df[col] = df[col].astype('category')

    end_mem = df.memory_usage().sum() / 1024**2
    print('Memory usage after optimization is: {:.2f} MB'.format(end_mem))
    print('Decreased by {:.1f}%'.format(100 * (start_mem - end_mem) / start_mem))
    
    return df

関数は、カグラーさんが作ってくれたもの(作成者が、kaggleから居なくなってしまった？)を使わせて頂きます。

本当にすべて必要か検討する

今回のデータセットは、1億レコードあるものでした。

データ全てを使用して訓練させる必要はあるのでしょうか？

よって、1億レコードの全てを使用する必要があるのかどうかを検討します。

1000万レコードほどで上位に行った人もいるので、全て使用する必要も無いと言えます。

と言っても、データが多いほど、柔軟性が高いモデルを作りやすく、スコアも上がりやすいのも事実なので、一概に言えないのも事実です。

分割してデータを渡して機械学習させる

機械学習アルゴリズムによっては、一度にデータを渡す物では無く、分割して渡すことができるものがあります(例：ニューラルネットワーク＝NN系)。

そんなときは、ファイルからある程度読み込むことを検討します。

以下のようなイメージです。

pandasの処理の高速化

pandasが、データ容量が多くなると処理遅くなりがちです。

特に、for文でデータの参照をするとものすごく遅くなってしまいます。

そこで、改善した方法を記載します。

改善方法は以下の通り

次に実装例を示します。

まとめ

最近は、データ容量が多いコンペが増えてきました。

今後も大量データの処理のスキルも必要になってくるといえます。