custom_tracer.py

import torch
from torch.fx import symbolic_trace, Tracer, Graph, GraphModule, Node
from typing import Any, Callable, Dict, Optional, Tuple, Union



"""
How to Create and Use Custom Tracers

`Tracer`--the class that implements the symbolic tracing functionality
of `torch.fx.symbolic_trace`--can be subclassed to override various
behaviors of the tracing process. In this tutorial, we'll demonstrate
how to customize the symbolic tracing process using some handwritten
Tracers. Each example will show that, by simply overriding a few methods
in the `Tracer` class, you can alter the Graph produced by symbolic
tracing. For a complete description of the methods that can be changed,
refer to the docstrings of the methods in the Tracer class. Information
can be found at: https://pytorch.org/docs/master/fx.html#torch.fx.Tracer

If you want a real-world example of a custom tracer, check out FX's AST
Rewriter in `rewriter.py`. `RewritingTracer` inherits from Tracer but
overrides the `trace` function so that we can rewrite all calls to
`assert` to the more FX-friendly `torch.assert`.

Note that a call to `symbolic_trace(m)` is equivalent to
`GraphModule(m, Tracer().trace(m))`. (`Tracer` is the default
implementation of Tracer as defined in `symbolic_trace.py`.)
"""



"""
Custom Tracer #1: Trace Through All `torch.nn.ReLU` Submodules

During symbolic tracing, some submodules are traced through and their
constituent ops are recorded; other submodules appear as an
atomic "call_module" Node in the IR. A module in this latter category
is called a "leaf module". By default, all modules in the PyTorch
standard library (`torch.nn`) are leaf modules. We can change this
by creating a custom Tracer and overriding `is_leaf_module`. In this
case, we'll keep the default behavior for all `torch.nn` Modules except
for `ReLU`.
"""

class M1(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        return self.relu(x)

default_traced: GraphModule = symbolic_trace(M1())
"""
Tracing with the default tracer and calling `print_tabular` produces:

    opcode       name    target    args       kwargs
    -----------  ------  --------  ---------  --------
    placeholder  x       x         ()         {}
    call_module  relu_1  relu      (x,)       {}
    output       output  output    (relu_1,)  {}

"""
default_traced.graph.print_tabular()

class LowerReluTracer(Tracer):
    def is_leaf_module(self, m : torch.nn.Module, qualname : str):
        if isinstance(m, torch.nn.ReLU):
            return False
        return super().is_leaf_module(m, qualname)

"""
Tracing with our custom tracer and calling `print_tabular` produces:

    opcode         name    target                             args       kwargs
    -------------  ------  ---------------------------------  ---------  ------------------
    placeholder    x       x                                  ()         {}
    call_function  relu_1  <function relu at 0x7f66f7170b80>  (x,)       {'inplace': False}
    output         output  output                             (relu_1,)  {}
"""
lower_relu_tracer = LowerReluTracer()
custom_traced_graph: Graph = lower_relu_tracer.trace(M1())
custom_traced_graph.print_tabular()



"""
Custom Tracer #2: Add an Extra Attribute to Each Node

Here, we'll override `create_node` so that we can add a new attribute to
each Node during its creation
"""

class M2(torch.nn.Module):
    def forward(self, a, b):
        return a + b

class TaggingTracer(Tracer):
    def create_node(self, kind : str, target : Union[str, Callable],
                    args : Tuple[Any], kwargs : Dict[str, Any], name : Optional[str] = None,
                    type_expr : Optional[Any] = None) -> Node:
        n = super().create_node(kind, target, args, kwargs, name)
        n.tag = "foo"
        return n

custom_traced_graph: Graph = TaggingTracer().trace(M2())

def assert_all_nodes_have_tags(g: Graph) -> bool:
    for n in g.nodes:
        if not hasattr(n, "tag") or not n.tag == "foo":
            return False
    return True

# Prints "True"
print(assert_all_nodes_have_tags(custom_traced_graph))