PyTorch で CUDA をサブプロセスで動かす

AIを複数プロセスで動かしたいことってあると思います。これまで、PyTorchを使って認識処理をサブプロセスで動かしたいなと思ったことがありましたが、以下のエラーが発生してあきらめていました。以下はCUDAを使う場合に発生します。

"Cannot re-initialize CUDA in forked subprocess. To use CUDA with "
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

原因

ちゃんとググってみると、以下の記事を見つけました。
https://stackoverflow.com/questions/48822463/how-to-use-pytorch-multiprocessing

内部の詳しいことは分からなかったのですが、要するに multiprocessing でサブプロセスを作成する時に fork システムコールで子プロセスを立ち上げると CUDA が対応できないらしいです。エラーメッセージに書かれている通りですね。。。

解決策

サブプロセスの生成方法を fork から spawn に変更すればよい。
また、multiprocessing の代わりに torch.multiprocessing を使う。

以下をサブプロセス生成前に記述する。

import torch.multiprocessing as multiprocessing
if multiprocessing.get_start_method() == 'fork':
    multiprocessing.set_start_method('spawn', force=True)
    print("{} setup done".format(multiprocessing.get_start_method()))