Pytorch 与 joblib 的 autograd 问题

将 pytorch 的 autograd 与 joblib 混合似乎存在问题。我需要并行获取大量样本的梯度。 Joblib 与 pytorch 的其他方面配合良好，但是，与 autograd 混合时会出现错误。我做了一个非常小的例子，显示串行版本工作正常，但并行版本崩溃。

from joblib import Parallel, delayed
import numpy as np
torch.autograd.set_detect_anomaly(True)
tt = lambda x, grad=True: torch.tensor(x, requires_grad=grad)

def Grad(X, Out):
    return autograd.grad(Out, X, create_graph=True, allow_unused=False)[0]

xs, ys = [], []
for i in range(10):
    xi = tt(np.random.rand()).float()    
    yi = xi * xi
    xs += [xi]
    ys += [yi]


Grads_serial = [Grad(x, y) for x, y in zip(xs, ys)]
print("Grads_serial", Grads_serial)
Grads_parallel = Parallel(n_jobs=2)([delayed(Grad)(x, y) for x, y in zip(xs, ys)])
print("Grads_parallel", Grads_parallel)

错误消息也不是很有帮助：

RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

Joblib 不会将与操作关联的图表复制到不同的进程。解决这个问题的一种方法是在作业内执行计算。

import torch
from torch import autograd
from joblib import Parallel, delayed
import numpy as np
torch.autograd.set_detect_anomaly(False)
tt = lambda x, grad=True: torch.tensor(x, requires_grad=grad)

def Grad(X, Out):
    # This will compute yi in the job, and thus will
    # create the graph here
    yi = Out[0](*Out[1])
    # now the differentiation works
    return autograd.grad(yi, X, create_graph=True, allow_unused=False)[0]

torch.set_num_threads(1)
xs, ys = [], []
for i in range(10):
    xi = tt(np.random.rand()).float()    
    yi = lambda xi: xi * xi, [xi]
    xs += [xi]
    ys += [yi]

Grads_serial = [Grad(x, y) for x, y in zip(xs, ys)]
print("Grads_serial", Grads_serial)
Grads_parallel = Parallel(n_jobs=2)([delayed(Grad)(x, y) for x, y in zip(xs, ys)])
print("Grads_parallel", Grads_parallel)

Edit

更多的哲学问题是

(1) 如果您可以简单地矢量化操作并让 torch 使用操作符内并行性，那么使用 joblib 并行性是否有意义？

(2) mak14 提到使用线程后端，它很好地修复了你的示例。但是多个线程将仅使用一个 CPU，这对于 IO 有限的作业有意义，例如发出 HTTP 请求，但对于 CPU 有限的操作则不然。

Edit #2

的存在torch.multiprocessing https://pytorch.org/docs/stable/notes/multiprocessing.html表明梯度需要一些特殊处理，您可以尝试使用以下命令将后端写入 joblibtorch.multiprocessing代替multiprocessing or threading.

在这里您可以找到有关如何在两个框架中构建图表的概述

https://www.tensorflow.org/guide/intro_to_graphs https://www.tensorflow.org/guide/intro_to_graphs

https://pytorch.org/blog/computational-graphs-constructed-in-pytorch/ https://pytorch.org/blog/computational-graphs-constructed-in-pytorch/

但我担心要给出一个明确的答案来解释为什么一个有效而另一个无效，就必须研究实施情况。