multiprocessingの備忘録 : 並列処理が動かなくなった時の対処法

機械学習モデルのパイプラインを作成していると、並列処理を行いたい場面が出てくるとおもいます。私も同様に、データの前処理→並列処理でモデル学習という流れでパイプラインを作成していました。しかし、前処理の内部実装を変えた途端に、突然並列処理が動かなくなりました。

結論

spawnを使う。

def main():
  preprocessing() # 前処理メソッド
  train_test() # 並列で学習するメソッド

if __name__=="__main__":
  multiprocessing.set_start_method('spawn') #ここを追加
  main()

環境

今まで正常に動いていたコード

今まで正常に動いていたコードは以下のようなものです。

import multiprocessing
import sklearn
import lightgbm as lgb

def preprocessing():
  # 前処理の内部実装

def train_fold(X_train, y_train, X_test, y_test):
  #lightgbmのハイパラを定義
  params=~
  X_train_for_lgb = lgb.Dataset(X_train, label=y_train)
  model=lgb.train(params, X_train_for_lgb)

  # テストデータで評価
  y_pred=model.predict(X_test)
  accuracy=sklearn.metrics.accuracy_score(y_test, y_pred)
  return accuracy


def train():
  # 並列で学習する内部実装
  #学習データXとラベルyを定義
  X,y=~~
  results=[]
  pool=multiprocessing.Pool(processes=4)
  kf=sklearn.model_selection.KFold(n_splits=4).split(X,y)
  for train_index, test_index in kf:
    results.append(pool.apply_async(train_fold, args=(X[train_index], y[train_index], X[test_index], y[test_index])))


def main():
  preprocessing() # 前処理メソッド
  train() # 並列で学習するメソッド

if __name__=="__main__":
  main()

前処理の内部実装を変えたら並列処理が動かなくなった

前処理メソッドの中で、sklearnのあるメソッドを使うようにしたら、突然並列処理が動かなくなりました。正確には、上記コードのtrain_fold()メソッドのX_train_for_lgbまでは実行されるのですが、lgb.train()以降が実行されなくなりました。

解決策と原因

冒頭にある通り、spawnを使うことで解決しました。

def main():
  preprocessing() # 前処理メソッド
  train_test() # 並列で学習するメソッド

if __name__=="__main__":
  multiprocessing.set_start_method('spawn') #ここを追加
  main()

原因は「メモリの使いすぎ」説が濃厚。 Ubuntuでは、デフォルトでforkを使っているようです。forkは、親プロセスのメモリを子プロセスにコピーするため、メモリを大量に消費します。一方、spawnは、親プロセスのメモリを子プロセスにコピーしないため、メモリを消費しないとのことです。 （しかし前処理と学習のコードはメソッドが分かれており、前処理が終了した段階でメモリが解放されるはずでは？という疑問は残ったままです。）

参考