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