CUDA multiprocess エラー

備忘録

Python上でCUDAを使ってGPUによるマルチプロセスを実行しようとすると下記のエラーに出てしまい、対処法を探すのにかなりの時間を費やした。

RuntimeError: CUDA error: initialization error


先人の知恵を借りながら、試行錯誤の結果、プログラムの先頭(Jupyter Notebookの一番初め)に下記を追加することでエラーが出なくなった。

import os

os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ['CUDA_LAUNCH_BLOCKING'] = "1"

Jupyter Notebook の場合は、一度カーネルを再起動し、上記の部分からやり直さないといけない。

前提として、nvidia-smiがエラーにならず、torch.cuda.is_available()がTrueになること。