reduce_mem_usageの高速化

Kaggleのテーブルコンペでreduce_mem_usageという関数を見かけます．
これはpandasのDataFrameのメモリ削減用の関数で，columnに含まれる数値の大きさに応じてpandasのデータ型(dtype)を変化させ，メモリ削減を行う関数です．特徴量生成後，学習を行う直前なんかに使われます．
簡単な具体的を挙げるとonehot特徴量 (0 or 1) のcolumnにint64は過剰でint8で十分だし，データ型を変換したらメモリが8分の1くらいに抑えられて嬉しいねみたいな話です．

ただ対象のDataFrameが大きい場合は，メモリ削減自体にかなり時間がかかってしまうので，今回はその関数の高速化に取り組みました．

目次

コード

以下提案する関数，変更箇所は後述

実装の中身は，for文を回してcolumnごとの最大・最小値と，各dtypeで表せる最大・最小値を比較して，値が変わらない範囲で出来るだけ使用メモリの小さいdtypeに変えています．

def reduce_mem_usage(df, verbose=True):
    numerics = ['int16', 'int32', 'int64', 'float16', 'float32', 'float64']
    start_mem = df.memory_usage().sum() / 1024**2 
    dfs = []
    for col in df.columns:
        col_type = df[col].dtypes
        if col_type in numerics:
            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:
                    dfs.append(df[col].astype(np.int8))
                elif c_min > np.iinfo(np.int16).min and c_max < np.iinfo(np.int16).max:
                    dfs.append(df[col].astype(np.int16))
                elif c_min > np.iinfo(np.int32).min and c_max < np.iinfo(np.int32).max:
                    dfs.append(df[col].astype(np.int32))
                elif c_min > np.iinfo(np.int64).min and c_max < np.iinfo(np.int64).max:
                    dfs.append(df[col].astype(np.int64) ) 
            else:
                if c_min > np.finfo(np.float16).min and c_max < np.finfo(np.float16).max:
                    dfs.append(df[col].astype(np.float16))
                elif c_min > np.finfo(np.float32).min and c_max < np.finfo(np.float32).max:
                    dfs.append(df[col].astype(np.float32))
                else:
                    dfs.append(df[col].astype(np.float64))
        else:
            dfs.append(df[col])
    
    df_out = pd.concat(dfs, axis=1)
    if verbose:
        end_mem = df_out.memory_usage().sum() / 1024**2
        num_reduction = str(100 * (start_mem - end_mem) / start_mem)
        print(f'Mem. usage decreased to {str(end_mem)[:3]}Mb:  {num_reduction[:2]}% reduction')
    return df_out

拾ったコードとの変更点

変えたのはデータ型を変換する部分で，デカいDataFrameを出来るだけ触らないように，データ型変換後のSeriesを都度listに入れて後からまとめてconcatするようにしました．

• [やり方1]見かけたコード*1

dtype変換して元のDataFrameに上書き

if c_min > np.iinfo(np.int8).min and c_max < np.iinfo(np.int8).max:
    df[col] = df[col].astype(np.int8)

• [やり方2]提案するコード

listにdtype変換後のSeriesを突っ込み，後でまとめてconcat

dfs = []
...

if c_min > np.iinfo(np.int8).min and c_max < np.iinfo(np.int8).max:
    dfs.append(df[col].astype(np.int8))

...
df_out = pd.concat(dfs, axis=1)

実験結果

行数:10000で固定し，列数を10, 100, 200で変化させてかかった時間を計測しました．

実行時間の比較

上書きするやり方は元のDataFrameが大きくなるほど実行時間が大きくなる一方で*2，concatするやり方は列数がこのぐらいの範囲だとあまり増加せずに済んでいます．
(実行時間は1回ずつしか回してないので多少の変動はあると思いますし，列数固定で行数を変更させたり，条件次第な部分はあると思います．)

特徴量生成前にメモリ削減する際の注意点(余談)

特徴量生成前にメモリ削減をすると望まないオーバーフローが起こる場合があるので注意が必要です．

簡単な実験とし， int8のDataFrameを作成し，2つのcolumnの和を取って8+120の計算をします．
int8のデータ型では-128~127までの値しか表現できないためやり方次第ではオーバーを起こします．

• DataFrameの作成

# DataFrameの作成, dtypeはint8
df = pd.DataFrame(data=[[8, 120]], dtype=np.int8)

print(df)
#    0    1
# 0  8  120

• [パターン1] sumで和をとる

print(df.sum(axis=1))
# 0    128
# dtype: int64

• [パターン2] df[col1] + df[col2]

print(df[0] + df[1])
# 0   -128
# dtype: int8

パターン2ではオーバーフローが起こり大小関係がおかしくなっているのでこのまま特徴量に使ったら精度の悪化に繋がりそうです．

特徴量生成前にメモリ削減関数のreduce_mem_usageを使うと特徴量生成の速度は上がりますが，オーバーフローが起こりうるので気にした方が良いかもしれません．

今回のコードはGithubに載せています．
実験コードも同じレポジトリ内に置いてます．
github.com

*1:見かけただの拾っただのという言い方をしているのは数多くのコンペで同じ関数を見かけていて，ある種共通のコード資産みたいになっているため，reduce_mem_usageと言えば共通認識を持てると判断した+初出がどれか分からないためです

*2:上書き時のindex周りの話だとは思いますが詳しい中身の話まではよく知りません