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 的向后兼容性。