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()