Python PyTorch rpc_async用法及代碼示例

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

用法:

torch.distributed.rpc.rpc_async(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,可選的) -用於此 RPC 的超時秒數。如果 RPC 在這段時間內沒有完成，則會引發一個表明它已超時的異常。值 0 表示無限超時，即永遠不會引發超時錯誤。如果未提供，則使用在初始化期間或使用_set_rpc_timeout 設置的默認值。

返回：

返回可以等待的 Future 對象。完成後，可以從 Future 對象中檢索args 和kwargs 上func 的返回值。

進行非阻塞 RPC 調用以在 worker to 上運行函數 func。 RPC 消息與 Python 代碼的執行並行發送和接收。此方法是線程安全的。此方法將立即返回一個可以等待的 Future 。

警告

不支持使用 GPU 張量作為參數或 func 的返回值，因為我們不支持通過網絡發送 GPU 張量。在將 GPU 張量用作參數或返回 func 值之前，您需要將它們顯式複製到 CPU。

警告

rpc_async API 不會複製參數張量的存儲，直到通過網絡發送它們，這可以由不同的線程完成，具體取決於 RPC 後端類型。調用者應確保這些張量的內容保持不變，直到返回的 Future 完成。

例子：

確保 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)
>>> fut1 = rpc.rpc_async("worker1", torch.add, args=(torch.ones(2), 3))
>>> fut2 = rpc.rpc_async("worker1", min, args=(1, 2))
>>> result = fut1.wait() + fut2.wait()
>>> 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)
>>> fut = rpc.rpc_async("worker1", my_script_add, args=(torch.ones(2), 3))
>>> ret = fut.wait()
>>> rpc.shutdown()

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