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