Python PyTorch remote用法及代碼示例

本文簡要介紹python語言中 torch.distributed.rpc.remote 的用法。

用法:

torch.distributed.rpc.remote(to, func, args=None, kwargs=None, timeout=- 1.0)

參數：

• to(str或者WorkerInfo或者int) -目標工作人員的名稱/等級/WorkerInfo。

• func(可調用的) -可調用函數，例如 Python 可調用函數、內置運算符(例如 add() )和帶注釋的 TorchScript 函數。

• args(tuple) -func 調用的參數元組。

• kwargs(dict) -是func 調用的關鍵字參數字典。

• timeout(float,可選的) -此遠程調用的超時(以秒為單位)。如果在此超時內未在此工作線程上成功處理 RRef 在工作線程 to 上的創建，則下次嘗試使用 RRef (例如 to_here() )時，將引發超時，指示這次失敗。值 0 表示無限超時，即永遠不會引發超時錯誤。如果未提供，則使用初始化期間或使用 _set_rpc_timeout 設置的默認值。

返回：

結果值的用戶 RRef 實例。使用阻塞 API torch.distributed.rpc.RRef.to_here() 在本地檢索結果值。

進行遠程調用以在工作線程 to 上運行 func，並立即將 RRef 返回到結果值。 Worker to 將是返回的 RRef 的所有者，而調用 remote 的 Worker 是一個用戶。所有者管理其 RRef 的全局引用計數，並且所有者 RRef 僅當全局沒有對其的活動引用時才被破壞。

警告

remote API 不會複製參數張量的存儲，直到通過線路發送它們，這可以由不同的線程完成，具體取決於 RPC 後端類型。調用者應確保這些張量的內容保持完整，直到返回的 RRef 得到所有者確認(可以使用 torch.distributed.rpc.RRef.confirmed_by_owner() API 進行檢查)。

警告

remote API 超時等錯誤將盡最大努力處理。這意味著當remote 發起的遠程調用失敗時，例如出現超時錯誤，我們會盡最大努力處理錯誤。這意味著錯誤會在異步基礎上處理並設置在生成的 RRef 上。如果在此處理之前應用程序沒有使用 RRef(例如 to_here 或 fork 調用)，那麽將來使用 RRef 將適當地引發錯誤。但是，用戶應用程序可能會在處理錯誤之前使用RRef。在這種情況下，可能不會引發錯誤，因為它們尚未被處理。

例子：

確保 MASTER_ADDR 和 MASTER_PORT 在兩個 worker 上都正確設置。有關詳細信息，請參閱 init_process_group() API。例如，

>>> export MASTER_ADDR=localhost
>>> export MASTER_PORT=5678

然後在兩個不同的進程中運行以下代碼：

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> rref1 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3))
>>> rref2 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 1))
>>> x = rref1.to_here() + rref2.to_here()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

下麵是使用 RPC 運行 TorchScript 函數的示例。

>>> # On both workers:
>>> @torch.jit.script
>>> def my_script_add(t1, t2):
>>>    return torch.add(t1, t2)

>>> # On worker 0:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> rref = rpc.remote("worker1", my_script_add, args=(torch.ones(2), 3))
>>> rref.to_here()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()