Python PyTorch backward用法及代码示例

本文简要介绍python语言中 torch.distributed.autograd.backward 的用法。

用法:

torch.distributed.autograd.backward(context_id: int, roots: List[Tensor], retain_graph=False) → None

参数：

• context_id(int) -我们应该为其检索梯度的 autograd 上下文 id。

• roots(list) -表示 autograd 计算根的张量。所有张量都应该是标量。

• retain_graph(bool,可选的) -如果为 False，则用于计算 grad 的图将被释放。请注意，几乎在所有情况下都不需要将此选项设置为 True，并且通常可以以更有效的方式解决。通常，您需要将其设置为 True 才能向后运行多次。

使用提供的根启动分布式反向传递。这目前实现了 FAST 模式算法，该算法假定在同一个分布式 autograd 上下文中跨工作人员发送的所有 RPC 消息将在后向传递期间成为 autograd 图的一部分。

我们使用提供的根来发现 autograd 图并计算适当的依赖关系。此方法阻塞，直到完成整个 autograd 计算。

我们在每个节点上适当的 torch.distributed.autograd.context 中累积梯度。根据调用 torch.distributed.autograd.backward() 时传入的 context_id，查找要使用的 autograd 上下文。如果没有与给定 ID 对应的有效自动分级上下文，我们会抛出错误。您可以使用get_gradients() API 检索累积的梯度。

例子：

>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     pred = model.forward()
>>>     loss = loss_func(pred, loss)
>>>     dist_autograd.backward(context_id, loss)