控制流 - Cond¶

torch.cond 是一個結構化的控制流運算元。它可用於指定類似 if-else 的控制流，並且從邏輯上可以被視為以下方式實現。

def cond(
    pred: Union[bool, torch.Tensor],
    true_fn: Callable,
    false_fn: Callable,
    operands: Tuple[torch.Tensor]
):
    if pred:
        return true_fn(*operands)
    else:
        return false_fn(*operands)

它獨特的優勢在於其表達數據依賴控制流的能力：它會降級為條件運算元 (torch.ops.higher_order.cond)，該運算元保留了謂詞、true 函式和 false 函式。這在編寫和部署模型時釋放了極大的靈活性，這些模型會根據張量運算的輸入或中間輸出的值或形狀來更改模型架構。

警告

torch.cond 是 PyTorch 中的原型功能。它對輸入和輸出類型的支持有限，目前不支持訓練。請期待未來版本的 PyTorch 中更穩定的實現。請在以下網址閱讀有關功能分類的更多資訊：https://pytorch.dev.org.tw/blog/pytorch-feature-classification-changes/#prototype

範例¶

以下是一個使用 cond 根據輸入形狀進行分支的範例

import torch

def true_fn(x: torch.Tensor):
    return x.cos() + x.sin()

def false_fn(x: torch.Tensor):
    return x.sin()

class DynamicShapeCondPredicate(torch.nn.Module):
    """
    A basic usage of cond based on dynamic shape predicate.
    """

    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        def true_fn(x: torch.Tensor):
            return x.cos()

        def false_fn(x: torch.Tensor):
            return x.sin()

        return torch.cond(x.shape[0] > 4, true_fn, false_fn, (x,))

dyn_shape_mod = DynamicShapeCondPredicate()

我們可以迫不及待地執行模型，並預期結果會根據輸入形狀而有所不同

inp = torch.randn(3)
inp2 = torch.randn(5)
assert torch.equal(dyn_shape_mod(inp), false_fn(inp))
assert torch.equal(dyn_shape_mod(inp2), true_fn(inp2))

我們可以匯出模型以進行進一步的轉換和部署

inp = torch.randn(4, 3)
dim_batch = torch.export.Dim("batch", min=2)
ep = torch.export.export(DynamicShapeCondPredicate(), (inp,), {}, dynamic_shapes={"x": {0: dim_batch}})
print(ep)

這會產生一個匯出的程式，如下所示

class GraphModule(torch.nn.Module):
    def forward(self, arg0_1: f32[s0, 3]):
        sym_size: Sym(s0) = torch.ops.aten.sym_size.int(arg0_1, 0)
        gt: Sym(s0 > 4) = sym_size > 4;  sym_size = None
        true_graph_0 = self.true_graph_0
        false_graph_0 = self.false_graph_0
        conditional: f32[s0, 3] = torch.ops.higher_order.cond(gt, true_graph_0, false_graph_0, [arg0_1]);  gt = true_graph_0 = false_graph_0 = arg0_1 = None
        return (conditional,)

    class <lambda>(torch.nn.Module):
        def forward(self, arg0_1: f32[s0, 3]):
            cos: f32[s0, 3] = torch.ops.aten.cos.default(arg0_1)
            sin: f32[s0, 3] = torch.ops.aten.sin.default(arg0_1);  arg0_1 = None
            add: f32[s0, 3] = torch.ops.aten.add.Tensor(cos, sin);  cos = sin = None
            return add

    class <lambda>(torch.nn.Module):
        def forward(self, arg0_1: f32[s0, 3]):
            sin: f32[s0, 3] = torch.ops.aten.sin.default(arg0_1);  arg0_1 = None
            return sin

請注意，torch.cond 已被降低到 torch.ops.higher_order.cond，它的 predicate 變成輸入形狀的 Symbolic 表達式，且分支函數變成頂層圖模組的兩個子圖屬性。

這裡有另一個範例，展示了如何表達依賴資料的控制流程

class DataDependentCondPredicate(torch.nn.Module):
    """
    A basic usage of cond based on data dependent predicate.
    """
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.cond(x.sum() > 4.0, true_fn, false_fn, (x,))

匯出後得到的匯出程式

class GraphModule(torch.nn.Module):
    def forward(self, arg0_1: f32[s0, 3]):
        sum_1: f32[] = torch.ops.aten.sum.default(arg0_1)
        gt: b8[] = torch.ops.aten.gt.Scalar(sum_1, 4.0);  sum_1 = None

        true_graph_0 = self.true_graph_0
        false_graph_0 = self.false_graph_0
        conditional: f32[s0, 3] = torch.ops.higher_order.cond(gt, true_graph_0, false_graph_0, [arg0_1]);  gt = true_graph_0 = false_graph_0 = arg0_1 = None
        return (conditional,)

    class <lambda>(torch.nn.Module):
        def forward(self, arg0_1: f32[s0, 3]):
            cos: f32[s0, 3] = torch.ops.aten.cos.default(arg0_1)
            sin: f32[s0, 3] = torch.ops.aten.sin.default(arg0_1);  arg0_1 = None
            add: f32[s0, 3] = torch.ops.aten.add.Tensor(cos, sin);  cos = sin = None
            return add

    class <lambda>(torch.nn.Module):
        def forward(self, arg0_1: f32[s0, 3]):
            sin: f32[s0, 3] = torch.ops.aten.sin.default(arg0_1);  arg0_1 = None
            return sin

torch.ops.higher_order.cond 的不變性¶

對於 torch.ops.higher_order.cond 有幾個有用的不變性

關於 predicate
- predicate 的動態性會被保留 (例如，上面範例中顯示的 gt)
- 如果 user-program 中的 predicate 是常數 (例如，一個 python bool 常數)，則運算子的 pred 將會是一個常數。
關於分支
- 輸入和輸出簽名將會是一個扁平化的 tuple。
- 它們是 torch.fx.GraphModule。
- 原始函數中的閉包會變成顯式的輸入。沒有閉包。
- 不允許對輸入或全域變數進行修改。
關於 operands
- 它也會是一個扁平化的 tuple。
user program 中 torch.cond 的巢狀結構會變成巢狀的圖模組。

API 參考¶

torch._higher_order_ops.cond.cond(pred, true_fn, false_fn, operands=())[source]¶

有條件地套用 true_fn 或 false_fn。

警告

torch.cond 是 PyTorch 中的原型功能。它對輸入和輸出類型的支持有限，目前不支持訓練。請期待未來版本的 PyTorch 中更穩定的實現。請在以下網址閱讀有關功能分類的更多資訊：https://pytorch.dev.org.tw/blog/pytorch-feature-classification-changes/#prototype

cond 是一個結構化的控制流程運算子。也就是說，它就像一個 Python 的 if 語句，但對 true_fn、false_fn 和 operands 有限制，使其可以使用 torch.compile 和 torch.export 進行擷取。

假設 cond 的參數滿足約束，cond 等同於以下內容

def cond(pred, true_branch, false_branch, operands):
    if pred:
        return true_branch(*operands)
    else:
        return false_branch(*operands)

參數

pred (Union[bool, torch.Tensor]) – 一個布林表達式或一個具有單一元素的 tensor，指示要套用哪個分支函數。
true_fn (Callable) – 一個可呼叫的函數 (a -> b)，它在被追蹤的範圍內。
false_fn (Callable) – 一個可呼叫的函數 (a -> b)，它在被追蹤的範圍內。true 分支和 false 分支必須具有一致的輸入和輸出，這表示輸入必須相同，並且輸出必須是相同的類型和形狀。
operands (Tuple of possibly nested dict/list/tuple of torch.Tensor) – true/false 函數的輸入 tuple。如果 true_fn/false_fn 不需要輸入，則可以為空。預設為 ()。

回傳類型

Any

範例

def true_fn(x: torch.Tensor):
    return x.cos()
def false_fn(x: torch.Tensor):
    return x.sin()
return cond(x.shape[0] > 4, true_fn, false_fn, (x,))

限制

條件陳述式 (aka pred) 必須滿足以下其中一項約束
- 它是一個只有一個元素的 torch.Tensor，且 dtype 為 torch.bool
- 它是一個布林表達式，例如 x.shape[0] > 10 或 x.dim() > 1 and x.shape[1] > 10
分支函數 (aka true_fn/false_fn) 必須滿足以下所有約束
- 函數簽名必須與 operands 相符。
- 函數必須回傳具有相同 metadata 的 tensor，例如形狀、dtype 等。
- 函數不能對輸入或全域變數進行 in-place 修改。(注意：允許分支中存在用於中間結果的 in-place tensor 運算，例如 add_)

警告

時間限制

分支的輸出必須是單一 Tensor。未來將支援 tensors 的 Pytree。

控制流 - Cond¶

範例¶

torch.ops.higher_order.cond 的不變性¶

API 參考¶

文件

教學

資源