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')とすることで期待通りの動作をする。