Python PyTorch wrap用法及代碼示例

本文簡要介紹python語言中 torch.fx.wrap 的用法。

用法:

torch.fx.wrap(fn_or_name)

參數：

fn_or_name(聯盟[str,可調用]) -調用時插入到圖中的全局函數的函數或名稱

可以在模塊級範圍調用此函數，將fn_or_name注冊為“leaf function”。 “leaf function” 將在 FX 跟蹤中保留為 CallFunction 節點，而不是通過以下方式進行跟蹤：

# foo/bar/baz.py
def my_custom_function(x, y):
    return x * x + y * y

torch.fx.wrap('my_custom_function')

def fn_to_be_traced(x, y):
    # When symbolic tracing, the below call to my_custom_function will be inserted into
    # the graph rather than tracing it.
    return my_custom_function(x, y)

此函數也可以等效地用作裝飾器：

# foo/bar/baz.py
@torch.fx.wrap
def my_custom_function(x, y):
    return x * x + y * y

包裝函數可以被認為是“leaf function”，類似於“leaf modules” 的概念，也就是說，它們是作為調用留在FX 跟蹤中而不是通過跟蹤的函數。

注意

保證此 API 的向後兼容性。