Pytorch 的 autograd 问题与 joblib

Question

将 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.

Answer 1

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)

编辑

更多的哲学问题是

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

(2) mak14 提到使用线程后端，它修复了您的示例是件好事。但是多个线程将只使用一个 CPU，这对 IO 有界作业有意义，比如发出 HTTP 请求，但对 CPU 有界操作没有意义。

编辑 #2

存在torch.multiprocessing建议梯度需要一些特殊处理，您可以尝试使用 torch.multiprocessing 而不是 multiprocessing 或 threading 将后端写入 joblib。

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

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

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

但我担心要给出一个明确的答案，为什么一个有效而另一个无效，则必须研究实现。