Python PyTorch Future.then用法及代码示例

本文简要介绍python语言中 torch.futures.Future.then 的用法。

用法:

then(callback)

参数：

callback(Callable) -Callable 将此 Future 作为唯一参数。

返回：

一个新的 Future 对象，它包含 callback 的返回值，并且在给定的 callback 完成时将被标记为已完成。

将给定的回调函数附加到此 Future ，该函数将在 Future 完成时运行。可以将多个回调添加到同一个 Future 中，但无法保证它们的执行顺序(要强制执行特定顺序，请考虑链接：fut.then(cb1).then(cb2) )。回调必须采用一个参数，即对此 Future 的引用。回调函数可以使用value()方法来获取值。请注意，如果此 Future 已完成，则给定的回调将立即内联运行。

如果 Future 的值包含驻留在 GPU 上的张量，则可能会在填充这些张量的异步内核尚未在设备上完成执行时调用回调。但是，回调将通过设置为当前(从全局池中获取)的一些专用流来调用，这些流将与这些内核同步。因此，回调对这些张量执行的任何操作都将在内核完成后安排在设备上。换句话说，只要回调不切换流，它就可以安全地操作结果，而无需任何额外的同步。这类似于 wait() 的非阻塞行为。

类似地，如果回调返回一个包含驻留在 GPU 上的张量的值，它可以这样做，即使产生这些张量的内核仍在设备上运行，只要回调在其执行期间没有更改流.如果想要更改流，则必须小心将它们与原始流重新同步，即调用回调时当前的流。

注意

请注意，如果回调函数通过异常完成并调用 fut.wait() 或通过回调中的其他代码抛出回调函数，则 then 返回的未来将被适当地标记为遇到的错误。但是，如果此回调稍后完成其他期货，则这些期货不会标记为已完成并出现错误，并且用户负责独立处理这些期货的完成/等待。

例子：

>>> def callback(fut):
...     print(f"RPC return value is {fut.wait()}.")
>>> fut = torch.futures.Future()
>>> # The inserted callback will print the return value when
>>> # receiving the response from "worker1"
>>> cb_fut = fut.then(callback)
>>> chain_cb_fut = cb_fut.then(
...     lambda x : print(f"Chained cb done. {x.wait()}")
... )
>>> fut.set_result(5)
RPC return value is 5.
Chained cb done. None