mp.set_start_method('spawn') 触发了一个错误，指出上下文已经设置

Question

这是我的完整代码

我已经成功地用一个小片段重现了我的主代码的行为。

在Google Colab Env 中，假设我设置hardware accelerator为GPU.

这是小片段：

import multiprocessing as mp

def foo(q):
    q.put('hello')

if __name__ == '__main__':
    mp.set_start_method('spawn')
    q = mp.Queue()
    p = mp.Process(target=foo, args=(q,))
    p.start()
    print(q.get())
    p.join()

使用mp.set_start_method('spawn')，我得到了错误RuntimeError: context has already been set，mp.set_start_method('spawn', force=True)如果我可以说，如果我使用它，它会陷入无限循环。

有什么办法可以防止在 Colab Env 中发生该错误？

PS请注意，mp.set_start_method('spawn')这行代码在我的代码中是必需的。否则，我收到错误RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

矫枉过正的解决方案

这里的问题是生成的子进程找不到__main__.foo.

一种（不太完美的）选择是将代码放在一个单独的文件中，例如创建一个新目录并将其添加到路径中。

我不想使用该解决方案，因为它对我的主代码来说太过分了。也许它会引导您找到更优雅的解决方案。

更新

这似乎是一个更合理的解决方案，但仍在寻找更优雅的答案。

Answer 1

您应该在函数调用（main 除外）之前将 start 方法设置为 spawn。

用法示例：

import multiprocessing as mp

try:
   mp.set_start_method('spawn', force=True)
   print("spawned")
except RuntimeError:
   pass

我通常使用这个块在 PyTorch 中进行多处理推理