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)