torch.func.functionalize

torch.func.functionalize(func, *, remove='mutations')

functionalize 是一種轉換，可用於從函式中移除（中間）變更和別名，同時保留函式的語意。

functionalize(func) 傳回一個與 func 具有相同語意的新函式，但已移除所有中間變更。在中間張量上執行的每個原地操作：intermediate.foo_() 都會被其異地等效項取代：intermediate_updated = intermediate.foo()。

functionalize 對於將 PyTorch 程式發送到無法輕易表示變更或別名運算子的後端或編譯器很有用。

參數

• func (Callable) – 一個接受一或多個參數的 Python 函式。

• remove (str) – 一個可選的字串參數，其值為 ‘mutations’ 或 ‘mutations_and_views’。 如果傳入 ‘mutations’，則所有變更運算子都將被替換為其非變更等效運算子。 如果傳入 ‘mutations_and_views’，則所有別名運算子也將被替換為其非別名等效運算子。 預設值：‘mutations’。

回傳值

回傳一個新的 “functionalized” 函式。 它接受與 func 相同的輸入，並具有相同的行為，但對函式中中間張量執行的任何變更（以及可選的別名）都將被移除。

回傳類型

Callable

functionalize 也會移除對函式輸入執行的變更（和檢視）。 然而，為了保留語意，functionalize 將在轉換完成後“修復”變更，方法是檢測是否有任何張量輸入“應該”被變更，並在必要時將新資料複製回輸入。

範例

>>> import torch
>>> from torch.fx.experimental.proxy_tensor import make_fx
>>> from torch.func import functionalize
>>>
>>> # A function that uses mutations and views, but only on intermediate tensors.
>>> def f(a):
...     b = a + 1
...     c = b.view(-1)
...     c.add_(1)
...     return b
...
>>> inpt = torch.randn(2)
>>>
>>> out1 = f(inpt)
>>> out2 = functionalize(f)(inpt)
>>>
>>> # semantics are the same (outputs are equivalent)
>>> print(torch.allclose(out1, out2))
True
>>>
>>> f_traced = make_fx(f)(inpt)
>>> f_no_mutations_traced = make_fx(functionalize(f))(inpt)
>>> f_no_mutations_and_views_traced = make_fx(functionalize(f, remove='mutations_and_views'))(inpt)
>>>
>>> print(f_traced.code)



def forward(self, a_1):
    add = torch.ops.aten.add(a_1, 1);  a_1 = None
    view = torch.ops.aten.view(add, [-1])
    add_ = torch.ops.aten.add_(view, 1);  view = None
    return add

>>> print(f_no_mutations_traced.code)



def forward(self, a_1):
    add = torch.ops.aten.add(a_1, 1);  a_1 = None
    view = torch.ops.aten.view(add, [-1]);  add = None
    add_1 = torch.ops.aten.add(view, 1);  view = None
    view_1 = torch.ops.aten.view(add_1, [2]);  add_1 = None
    return view_1

>>> print(f_no_mutations_and_views_traced.code)



def forward(self, a_1):
    add = torch.ops.aten.add(a_1, 1);  a_1 = None
    view_copy = torch.ops.aten.view_copy(add, [-1]);  add = None
    add_1 = torch.ops.aten.add(view_copy, 1);  view_copy = None
    view_copy_1 = torch.ops.aten.view_copy(add_1, [2]);  add_1 = None
    return view_copy_1


>>> # A function that mutates its input tensor
>>> def f(a):
...     b = a.view(-1)
...     b.add_(1)
...     return a
...
>>> f_no_mutations_and_views_traced = make_fx(functionalize(f, remove='mutations_and_views'))(inpt)
>>> #
>>> # All mutations and views have been removed,
>>> # but there is an extra copy_ in the graph to correctly apply the mutation to the input
>>> # after the function has completed.
>>> print(f_no_mutations_and_views_traced.code)



def forward(self, a_1):
    view_copy = torch.ops.aten.view_copy(a_1, [-1])
    add = torch.ops.aten.add(view_copy, 1);  view_copy = None
    view_copy_1 = torch.ops.aten.view_copy(add, [2]);  add = None
    copy_ = torch.ops.aten.copy_(a_1, view_copy_1);  a_1 = None
    return view_copy_1

functionalize 存在一些值得指出的“失敗模式”

1. 與其他 torch.func 轉換一樣，functionalize() 不適用於直接使用 .backward() 的函式。 對於 torch.autograd.grad 也是如此。 如果您想使用 autograd，您可以直接使用 functionalize(grad(f)) 計算梯度。

2. 與其他 torch.func 轉換一樣，functionalize() 不適用於全域狀態。 如果您在取得非本機狀態檢視/變更的函式上呼叫 functionalize(f)，functionalization 將會簡單地執行 no-op 並將檢視/變更呼叫直接傳遞到後端。 一種解決方法是確保將任何非本機狀態建立包裝到一個更大的函式中，然後您再對其呼叫 functionalize。

3. resize_() 有一些限制：functionalize 只能在程式中使用 resize_()` 的情況下運作，前提是正在調整大小的張量不是檢視。

4. as_strided() 有一些限制：functionalize 無法處理產生具有重疊記憶體的張量的 as_strided() 呼叫。

最後，一個有助於理解 functionalization 的心智模型是，大多數使用者 PyTorch 程式都是使用公開的 torch API 編寫的。 執行時，torch 運算子通常會被分解為我們的內部 C++ “ATen” API。 functionalization 的邏輯完全在 ATen 層級發生。 Functionalization 知道如何取得 ATen 中的每個別名運算子，並將其對應到其非別名等效運算子（例如 tensor.view({-1}) -> at::view_copy(tensor, {-1})），以及如何取得 ATen 中的每個變更運算子，並將其對應到其非變更等效運算子（例如 tensor.add_(1) -> at::add(tensor, -1)），同時追蹤別名和變更以了解何時進行修復。 有關哪些 ATen 運算子是別名或變更運算子的所有資訊都來自 https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml。