Pandasで大きなファイルを読み込む

機械学習

Kaggleの勉強のために「Kaggleコンペティション チャレンジブック」を読み始めた。最初に扱う題材は santander-product-recommendation だった。

Santander Product Recommendation
Can you pair products with people?

早速、Kaggle Notebookで訓練データを読み込もうとするとメモリ不足でエラーになった。試しに、GoogleのColaboratoryで実行すると正常に終了した。読み込んだデータの容量を確認すると5GB程度だった。

import pandas as pd
import numpy as np

trn = pd.read_csv('train_ver2.csv')
print(trn.memory_usage(index=True).sum())

# /usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py:2718: DtypeWarning: Columns (5,8,11,15) have mixed types.Specify dtype option on import or set low_memory=False.
#  interactivity=interactivity, compiler=compiler, result=result)
# 5240566784

2021/12/14時点でKaggleのNotebookは無料版でRAMが16GB割り当てられているので問題なく読み込めそうな容量だったのが不思議。Pandasのバージョンの影響か確認するためにColabにインストールされているpandas==1.1.5をインストールしても読み込めなかった。

Colabで作業はできるが、Kaggle Notebookの方がメモリが多いのでできればなんとかしたい。調べてみると、Pandasの小技があるようだったので試す。

型を指定する

Pandasでread_csvをすると、より大きな型で値が保存されるようなので、事前にどのようなデータかわかればメモリ容量を削減できる。

thead = pd.read_csv('/kaggle/input/santander-product-recommendation/train_ver2.csv.zip', usecols=[47])
print(thead.memory_usage(index=True).sum())
print(thead.info())

# 109178600
# <class 'pandas.core.frame.DataFrame'>
# RangeIndex: 13647309 entries, 0 to 13647308
# Data columns (total 1 columns):
#  #   Column           Dtype
# ---  ------           -----
#  0   ind_recibo_ult1  int64
# dtypes: int64(1)
# memory usage: 104.1 MB
# None
# array([0, 1])

48個目の値を読み込むと100MB消費しているのがわかる。値は0,1しかないがint64の型に入っているので、これをbool型になるようにする。

thead = pd.read_csv('/kaggle/input/santander-product-recommendation/train_ver2.csv.zip', usecols=[47], dtype={'ind_recibo_ult1': 'bool'})
print(thead.memory_usage(index=True).sum())
print(thead.info())
thead['ind_recibo_ult1'].unique()

# 13647437
# <class 'pandas.core.frame.DataFrame'>
# RangeIndex: 13647309 entries, 0 to 13647308
# Data columns (total 1 columns):
#  #   Column           Dtype
# ---  ------           -----
#  0   ind_recibo_ult1  bool 
# dtypes: bool(1)
# memory usage: 13.0 MB
# None
# array([False,  True])

boolとして読み込むと13MBまで節約できた。他のカラムも同様に対応する。

# カラム名を取得するために一部だけ読み込む
thead = pd.read_csv('/kaggle/input/santander-product-recommendation/train_ver2.csv.zip', nrows=5)

# 24個目以降のデータは0,1
dtypes = dict(zip(thead.columns.values[24:], ['bool' for i in range(24)]))
# 次の二つは nan が含まれるので object
dtypes['ind_nomina_ult1'] = 'object'
dtypes['ind_nom_pens_ult1'] = 'object'
train = pd.read_csv('/kaggle/input/santander-product-recommendation/train_ver2.csv.zip', dtype=dtypes)
print(train.memory_usage(index=True).sum())

# 3138881198

全てのデータを読み込んでも3GB程度になりメモリに乗せることができた。

コメント