torch.multiprocessingでfork時に親プロセスでmodel()した後に子プロセスでmodel()できない
自分のプロセス周りの知識不足のせいだと思うが、納得のいかない挙動だったのでログ。
import numpy as np import torch from torch.autograd import Variable import torch.multiprocessing as mp from model import NN def f(x): model = NNPolicy(1, 6) print(model.forward(x)) if __name__ == '__main__': x = Variable(torch.Tensor(np.random.randn(1, 1, 80, 80))) f(x) processes = [] for i in range(5): p = mp.Process(target=f, args=(x,)) p.start() processes.append(p) for p in processes: p.join()
NNはpytorchの一般的なNNモデルでここでは問題の本質ではないと思い省略。
このコードは以下のような出力になる。
tensor([[-0.0724]], grad_fn=<AddmmBackward>)
これはおかしい挙動で、本来なら以下のように親プロセスのprintと子プロセスのプリントx5で計6回のprintが実行されるはずである。
tensor([[-0.0724]], grad_fn=<AddmmBackward>) tensor([[-0.0724]], grad_fn=<AddmmBackward>) tensor([[-0.0724]], grad_fn=<AddmmBackward>) tensor([[-0.0724]], grad_fn=<AddmmBackward>) tensor([[-0.0724]], grad_fn=<AddmmBackward>) tensor([[-0.0724]], grad_fn=<AddmmBackward>)
ここで、親プロセスの f(x)
をコメントアウトすると、5回のprint出力が得られる。
これは、親プロセスのmodelを子プロセスに渡してもxを子プロセスで定義しても起きる。
親プロセスの f(x)
を print(model.forward(x))
に変えても起きる。
必ず親プロセスでmodel()を実行すると子プロセスのmodel()が実行されない。エラーも投げられない。
調べてみるとPython(UNIX)のforkはcopy-on-writeらしいので、その辺が原因か?んー。
PythonにおけるShared Memory - KZKY memo
ちなみに、上記のコードは __main__
以下にmp.set_start_method('spawn')
とすることで期待通りの動作をする。