92 lines
3.7 KiB
Python
92 lines
3.7 KiB
Python
|
# type: ignore
|
|||
|
r'''
|
|||
|
**This feature is under a Beta release and its API may change.**
|
|||
|
|
|||
|
FX is a toolkit for developers to use to transform ``nn.Module``
|
|||
|
instances. FX consists of three main components: a **symbolic tracer,**
|
|||
|
an **intermediate representation**, and **Python code generation**. A
|
|||
|
demonstration of these components in action:
|
|||
|
|
|||
|
::
|
|||
|
|
|||
|
import torch
|
|||
|
# Simple module for demonstration
|
|||
|
class MyModule(torch.nn.Module):
|
|||
|
def __init__(self):
|
|||
|
super().__init__()
|
|||
|
self.param = torch.nn.Parameter(torch.rand(3, 4))
|
|||
|
self.linear = torch.nn.Linear(4, 5)
|
|||
|
|
|||
|
def forward(self, x):
|
|||
|
return self.linear(x + self.param).clamp(min=0.0, max=1.0)
|
|||
|
|
|||
|
module = MyModule()
|
|||
|
|
|||
|
from torch.fx import symbolic_trace
|
|||
|
# Symbolic tracing frontend - captures the semantics of the module
|
|||
|
symbolic_traced : torch.fx.GraphModule = symbolic_trace(module)
|
|||
|
|
|||
|
# High-level intermediate representation (IR) - Graph representation
|
|||
|
print(symbolic_traced.graph)
|
|||
|
"""
|
|||
|
graph(x):
|
|||
|
%param : [#users=1] = self.param
|
|||
|
%add_1 : [#users=1] = call_function[target=<built-in function add>](args = (%x, %param), kwargs = {})
|
|||
|
%linear_1 : [#users=1] = call_module[target=linear](args = (%add_1,), kwargs = {})
|
|||
|
%clamp_1 : [#users=1] = call_method[target=clamp](args = (%linear_1,), kwargs = {min: 0.0, max: 1.0})
|
|||
|
return clamp_1
|
|||
|
"""
|
|||
|
|
|||
|
# Code generation - valid Python code
|
|||
|
print(symbolic_traced.code)
|
|||
|
"""
|
|||
|
def forward(self, x):
|
|||
|
param = self.param
|
|||
|
add_1 = x + param; x = param = None
|
|||
|
linear_1 = self.linear(add_1); add_1 = None
|
|||
|
clamp_1 = linear_1.clamp(min = 0.0, max = 1.0); linear_1 = None
|
|||
|
return clamp_1
|
|||
|
"""
|
|||
|
|
|||
|
The **symbolic tracer** performs “symbolic execution” of the Python
|
|||
|
code. It feeds fake values, called Proxies, through the code. Operations
|
|||
|
on theses Proxies are recorded. More information about symbolic tracing
|
|||
|
can be found in the :func:`symbolic_trace` and :class:`Tracer`
|
|||
|
documentation.
|
|||
|
|
|||
|
The **intermediate representation** is the container for the operations
|
|||
|
that were recorded during symbolic tracing. It consists of a list of
|
|||
|
Nodes that represent function inputs, callsites (to functions, methods,
|
|||
|
or :class:`torch.nn.Module` instances), and return values. More information
|
|||
|
about the IR can be found in the documentation for :class:`Graph`. The
|
|||
|
IR is the format on which transformations are applied.
|
|||
|
|
|||
|
**Python code generation** is what makes FX a Python-to-Python (or
|
|||
|
Module-to-Module) transformation toolkit. For each Graph IR, we can
|
|||
|
create valid Python code matching the Graph’s semantics. This
|
|||
|
functionality is wrapped up in :class:`GraphModule`, which is a
|
|||
|
:class:`torch.nn.Module` instance that holds a :class:`Graph` as well as a
|
|||
|
``forward`` method generated from the Graph.
|
|||
|
|
|||
|
Taken together, this pipeline of components (symbolic tracing →
|
|||
|
intermediate representation → transforms → Python code generation)
|
|||
|
constitutes the Python-to-Python transformation pipeline of FX. In
|
|||
|
addition, these components can be used separately. For example,
|
|||
|
symbolic tracing can be used in isolation to capture a form of
|
|||
|
the code for analysis (and not transformation) purposes. Code
|
|||
|
generation can be used for programmatically generating models, for
|
|||
|
example from a config file. There are many uses for FX!
|
|||
|
|
|||
|
Several example transformations can be found at the
|
|||
|
`examples <https://github.com/pytorch/examples/tree/master/fx>`__
|
|||
|
repository.
|
|||
|
'''
|
|||
|
|
|||
|
from .graph_module import GraphModule
|
|||
|
from .symbolic_trace import symbolic_trace, Tracer, wrap
|
|||
|
from .graph import Graph
|
|||
|
from .node import Node, map_arg
|
|||
|
from .proxy import Proxy
|
|||
|
from .interpreter import Interpreter as Interpreter, Transformer as Transformer
|
|||
|
from .subgraph_rewriter import replace_pattern
|