Pythonで中程度〜重めのデータを読み込むときの高速化テク

バージョン: Python 3.6以降ならたぶんなんでもOK

これを使えば、十数秒〜数分かかるデータのロードを、1~2秒という短い時間で終わらせることができます（だいたいできる）。

ざっと説明すると、一度ロードしたデータを pickle というモジュールで pkl 形式で保存し2回目以降の読み込みは pkl のバイナリからロードすることで高速化する、というシンプルな方法です。これに加え、Pythonのデコレータ記法を加えることで、コードの責任を分離しつつ可読性の低下を抑えます。

まとめ

先にまとめておきます

pickle_hook() という関数（「説明」にて後述）をコピペする
データのロード関数の前行に @pickle_hook と追記
pkl を保存する場所の指定（関数の引数と呼び出しに path_pkl を追加）

準備に対して時短が高効率であり、めちゃくちゃコスパが良いです。

説明

Python には pickle という標準モジュールがあります。 pickle はあらゆるPythonオブジェクトを格納、バイナリ化しpkl としてファイルを作成・ロードできるモジュールです。 pip 等で追加インストール不要で、import pickle でインポート可能です。

似たようなものに TFRecord がありますが、TPUを使わない場合、個人的にやや使いにくいと思っています。なのでいつもこの方法（pickle化）で実験をしたりコンペデータを扱っています。

Pythonにはデコレータと呼ばれる関数を修飾する記法があります（@から始まるアレ）。これを利用し、関数のアウトプットをキャッチし、pickle化します。

以下がpickle化の関数 pickle_hook() です。これをロード部分にデコレートすることになります。やり方としては、データをロードする関数に対し pkl ファイルを生成する関数をデコレートするだけです。

f:id:hyper-pigeon:20210805153252p:plain — スマホ用コード（画像にしてあります）

ブログ閲覧者のスマホユーザ率が50%を超えているので、ここでは画像を掲載しています。テキストベースのソースコードも本記事の最後にも用意してあります。
@wraps 部分は、多重デコレータに対応するためなので、他のデコレータと被っていなければ必要ないです。なので実質的なコードは十数行です。

例

雰囲気を伝えたくてコードを載せていますが、そこまでちゃんと読む必要はないです（重要でない部分は...で省略しています）。 要するにロード部分の関数の前行に @pickle_hook を追加、また関数の呼び出しに保存名を指定することで発動します。 一度実行されると、ファイルが生成され、次回からロード部分の関数をスキップし pkl ファイルを読みに行くようになります。

このあいだのatmaCupのときのデータをロードする部分です。ハードのI/Oにも依存しますが、atmaCup11の場合Colabで1~2秒でロード可能です。

f:id:hyper-pigeon:20210805070452p:plain — スマホ用コード（画像にしてあります）

これにより、「学習データの画像リスト」と「学習画像の名前のリスト」、「テスト画像のリスト」と「テスト画像の名前のリスト」の4オブジェクトをタプル化したPythonオブジェクトが保存されることになります。

画像をテンソル化せずにリスト化しているのは、画像解像度が一枚一枚異なるデータだったからです。後からクロップや縮小ができるような配慮です。

メリット・デメリット

メリット

読み込みがメチャクチャ早い
大量のデータを filename.pkl のように1つのデータとして扱うことができ、圧縮解凍などの手間がなく、ファイルの移動が便利（scp にしろ、GoodleDrive/Kaggleにアップロードにしろ）
- たとえばローカルで作成し、それを kaggle Notebook や、Google Colaboratory に移植することで、データセットを丸々移動せず擬似的にデータを扱えるようになる

デメリット

作成時、あまりにも重いデータの場合メモリがオーバーしクラッシュする（対処方法あり）。またメモリ全載せ前提なのでデータがデカすぎる場合成立しない。
コードの変更と対応できていないファイルが作成される（使い捨てのコンペに向いています）

向いているもの・向いてないもの

向いているもの

画像データ
ファイル数の多いデータ（並サイズの画像だと1,000データ程度以上からアドバンテージが生まれます）
前処理に時間のかかるテーブルデータなど

向いてないもの

重すぎるデータ: RAMの半分ギガバイト（RAMが24GBなら12GB）を超えるデータ / そもそもメモリに載らないデータ
軽量データ（恩恵が少ないです）
頻繁に変更が発生するデータ

コツとか気をつける点

float32/64で画像を保存しないこと
- 絶対ダメです。よほどの意味がない限りは uint8 で保存すべきです。float64 で保存すると、単純計算でuint8の8倍の容量を食い、保存/ロード速度・メモリ効率が落ち、pickle の恩恵が得られず、遅くなる場合が多いです。なので読み込んだあとで x = x.astype(np.float32)/255.0 みたいに変換するのがいいでしょう
pkl ファイル作成中にエラーが出たりするとバグファイルが残るので、そのファイルは毎回消去しましょう。できれば実行が確認できたコードに対してデコレータを施すのがいいでしょう
あまりにも重い単一の pkl ファイルを作らないこと（メモリと相談）
- 一応適度なサイズに pkl を分割して作成するコードもあるので、需要があればまたブログにします

おわりに

この方法の真髄は、取り回しの良さにあります。しかもモジュール化されており、わかりやすく、あらゆるオブジェクトを扱えるので自由度が高いです。

もし記事の反響がいい感じだった場合、分割保存と併せてPyPI化します。（はてなブクマ数5くらい?）

PC向けコード（テキストベース）

pickle_hook本体

import os
import pickle
from functools import wraps

def pickle_hook(func):
    """decorator for dataloader, load pkl if once loaded.
    USAGE:
        @pickle_hook
        loadfunc(path_pkl):
            <procedure>
            return data
        data = loadfunc(path_pkl="path to .pkl")
    """
    @wraps(func)
    def wrapper(*args, **kwargs):

        path_pkl = kwargs["path_pkl"]

        if(os.path.exists(path_pkl) == True):
            with open(path_pkl, mode='rb') as f:
                X = pickle.load(f)
            return X
        else:
            X = func(*args, **kwargs)
            with open(path_pkl, mode='wb') as f:
                pickle.dump(X, f, protocol=4)
            return X

    return wrapper

例

from scripts.func_gp import pickle_hook

# ...

class DataHolder:
    def __init__(self, get_by_list=False):

# ...

    def load_data(self):
        print('[loading images]')
        self.train_img, self.test_img, self.train_name, self.test_name = self.read_images(path_pkl=os.path.join(self.path_pkl, 'images.pkl'))

# ...

    @pickle_hook
    def read_images(self, path_pkl):
        train_img, test_img = [], []
        train_name, test_name = [], []
        
        # train
        print('[load] train')
        for i in range(len(self.df_train['object_id'])):
            file_id = self.df_train['object_id'][i]
            train_img.append(self.read_image(file_id))
            train_name.append(file_id)
            
        # test
        print('[load] test')
        for i in range(len(self.df_test['object_id'])):
            file_id = self.df_test['object_id'][i]
            test_img.append(self.read_image(file_id))
            test_name.append(file_id)

        return (train_img, test_img, train_name, test_name)

CC56

CC56: 567収束までに達成したい56リスト