pytorch 收集失败 sparse_grad=True

Question

即使是非常简单的例子，如果sparse_grad=True，backward()也无法工作，请看下面的错误。

这个错误是预料之中的，还是我使用 gather 的方式不对？

In [1]: import torch as th

In [2]: x = th.rand((3,3), requires_grad=True)

# sparse_grad = False, the backward could work as expetecd
In [3]: th.gather(x @ x, 1, th.LongTensor([[0], [1]]), sparse_grad=False).sum().backward()

# sparse_grad = True, backward CANNOT work
In [4]: th.gather(x @ x, 1, th.LongTensor([[0], [1]]), sparse_grad=True).sum().backward()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
----> 1 th.gather(x @ x, 1, th.LongTensor([[0], [1]]), sparse_grad=True).sum().backward()

~/miniconda3/lib/python3.9/site-packages/torch/_tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)
    305                 create_graph=create_graph,
    306                 inputs=inputs)
--> 307         torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
    308
    309     def register_hook(self, hook):

~/miniconda3/lib/python3.9/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    152         retain_graph = create_graph
    153
--> 154     Variable._execution_engine.run_backward(
    155         tensors, grad_tensors_, retain_graph, create_graph, inputs,
    156         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag

RuntimeError: sparse tensors do not have strides

Answer 1

我认为torch.gather不支持稀疏运算符：

torch.gather(x, 1, torch.LongTensor([[0], [1]]).to_sparse())

结果：

NotImplementedError: Could not run 'aten::gather.out' with arguments from the 'SparseCPU' backend.

我认为你应该在 pytorch's github 上提出问题或功能请求。

Answer 2

torch.gather 确实支持稀疏操作，因此您可以设置关键参数 sparse_grad=True。请参阅文档：https://pytorch.org/docs/stable/generated/torch.gather.html

但是，如果您确实使用稀疏梯度，则可以使用的操作数会减少。有些操作根本没有为稀疏梯度定义，其他操作没有实现。

从你的错误来看，我认为你在某处使用了内核或类似的操作，需要你设置一个步幅。该步幅需要一个完整的梯度，因为它是为一个完整的梯度定义的。由于你只有一个稀疏梯度，stride 会尝试将梯度作为在稀疏操作中尚未计算的输入。

如果你想让你的代码使用稀疏梯度，你必须将损失计算限制在稀疏梯度上。不要使用步幅或任何其他需要完整渐变的东西。

pytorch 收集失败 sparse_grad=True

pytorch gather failed with sparse_grad=True

python

pytorch

backpropagation