Python PyTorch Transformer用法及代碼示例

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

用法:

class torch.fx.Transformer(module)

參數：

module(GraphModule) -要轉換的Module。

Transformer 是一種特殊類型的解釋器，它產生一個新的 Module 。它公開了一個 transform() 方法，該方法返回轉換後的 Module 。 Transformer 不需要參數來運行，就像 Interpreter 一樣。 Transformer 完全象征性地工作。

示例

假設我們想用torch.sigmoid 交換所有torch.neg 實例，反之亦然(包括它們的Tensor 方法等價物)。我們可以像這樣子類化Transformer：

class NegSigmSwapXformer(Transformer):
    def call_function(self, target : 'Target', args : Tuple[Argument, ...], kwargs : Dict[str, Any]) -> Any:
        if target == torch.sigmoid:
            return torch.neg(*args, **kwargs)
        return super().call_function(n)

    def call_method(self, target : 'Target', args : Tuple[Argument, ...], kwargs : Dict[str, Any]) -> Any:
        if target == 'neg':
            call_self, *args_tail = args
            return call_self.sigmoid(*args_tail, **kwargs)
        return super().call_method(n)

def fn(x):
    return torch.sigmoid(x).neg()

gm = torch.fx.symbolic_trace(fn)

transformed : torch.nn.Module = NegSigmSwapXformer(gm).transform()
input = torch.randn(3, 4)
torch.testing.assert_allclose(transformed(input), torch.neg(input).sigmoid())

注意

保證此 API 的向後兼容性。