Python PyTorch rpc_sync用法及代碼示例

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

用法:

torch.distributed.rpc.rpc_sync(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 設置的默認值。

返回：

返回使用 args 和 kwargs 運行 func 的結果。

進行阻塞 RPC 調用以在 worker to 上運行函數 func。 RPC 消息與 Python 代碼的執行並行發送和接收。此方法是線程安全的。

例子：

確保 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)
>>> ret = rpc.rpc_sync("worker1", torch.add, args=(torch.ones(2), 3))
>>> 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)
>>> ret = rpc.rpc_sync("worker1", my_script_add, args=(torch.ones(2), 3))
>>> rpc.shutdown()

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