sklearnやkearsのCross Validationを並列で実行して高速化する

Kaggleに取り組んでいた時、実行時間制限付きのノートブックコンペがあり、モデルを高速に実行する必要が出てきた。

そこで、クロスバリデーションなら同時に更新する値や配列が無いため、なんとか並列実行出来ないか調べてみたところ、joblibを使ってparallelに実行して高速化することが出来た。

joblib.Parallel — joblib 1.5.dev0 documentation

exampleを見てみる。

>>> from math import sqrt
>>> from joblib import Parallel, delayed
>>> Parallel(n_jobs=1)(delayed(sqrt)(i**2) for i in range(10))
[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]

Parallel（並列実行数）（delayed（関数名）（引数）イテレータ））という並びで記述すればいいっぽい。

訓練データをtrain、target、テストデータをtest、使いたい分類機をMyModelとして、KFoldをParallelにしてみる。。

oof = np.zeros(len(train))
pred = np.zeros(len(test))

kf = StratifiedKFold(n_splits=NFOLDS, shuffle=True, random_state=42)


def one_fold(fold, trn_idx, val_idx):
    print('fold', fold)
    X_train, X_val = X.loc[trn_idx, :], X.loc[val_idx, :]
    y_train, y_val = y.loc[trn_idx, :], y.loc[val_idx, :]

    model = MyModel()
    model.fit(X_train, y_train)
    _oof = model.predict(X_val)
    _pred = model.predixt(test)

    return [val_idx, _oof, _pred]


result = joblib.Parallel(n_jobs=-1, verbose=0)(
    joblib.delayed(one_fold)(fold, trn_idx, val_idx)
    for fold, (trn_idx, val_idx) in enumerate(kf.split(train, target)))

for val_idx, _oof, _pred in result:
    oof[val_idx] = _oof
    pred += _pred / NFOLDS

並列に実行されるので、printされるfoldの順番はばらばらになる。

fold 4
fold 1
fold 0
fold 3
fold 2

あとはMyModelのところをkerasなりxgboostなりlightgbmなりrandom forestなり好きなモデルに置き換えればOK。

GPUを使う場合、並列にする分だけのVRAMが無いと実行できないので注意する。

全結合だけのニューラルネットやxgboostであればGPUで並列に実行すると高速化が実感できるが、CNNはGPUのコアを使い切るみたいで並列にしても高速化されなかった。

モデルに合わせて高速化されるかどうか違うようだ。

また、重みを保存する場合は save_name + fold みたいにして各 fold で違う名前にしておく。

同じ名前で保存するようなプログラムだと、並列に実行される途中で保存作業が他と被って正しく保存されない。

参考

Attention Required! | Cloudflare