tracer — Model Optimizer 0.29.0 (original) (raw)
Tracing module for building a set of partial Graphs describing the model.
Classes
| RobustTracer | A robust graph tracer for tracing dependencies between layers. |
|---|---|
| GraphCollection | A collection of fx graphs and tracing results. |
Functions
| recursive_trace | Recursively trace the model and all failing sub-modules and return graph collection. |
|---|
class GraphCollection
Bases: object
A collection of fx graphs and tracing results.
__init__()
Initialize the graph collection.
Return type:
None
failure_msg(m)
Get the failure message of the module if it failed.
Parameters:
m (Module | Callable) –
Return type:
str | None
is_failed(m)
Check if node target failed during tracing.
Parameters:
m (Module | Callable) –
Return type:
bool
is_unvisited(m)
Check if the node target is unvisited.
Parameters:
m (Module | Callable) –
Return type:
bool
recursive_trace(model, concrete_args=None)
Recursively trace the model and all failing sub-modules and store graph collection.
Parameters:
- model (Module) – The model to be traced.
- concrete_args (dict [ str , Any ] | None) – [EXPERIMENTAL; not backwards compatible!] The concrete arguments to be used for tracing of the root module (i.e. the model’s forward method). Note that if/when the tracing recurses, no concrete args are passed on! Defaults to None.
Return type:
None
Note that the graph collection contains the graph of the root module as well as all failing submodules. For failing modules, the graph will contain the failing module as a leaf and then separate graphs for each submodule within the failing module to maximize the coverage of the tracing results.
class RobustTracer
Bases: object
A robust graph tracer for tracing dependencies between layers.
Note that this tracer wraps unsupported ops into its parent module as a whole.
__init__()
Initialize the tracer.
failure_msg(m)
Get the failure message of the module.
Parameters:
m (Module | Callable) –
Return type:
str | None
is_failed(m)
Check if the module is failed.
Parameters:
m (Module | Callable) –
Return type:
bool
is_leaf_module(m, module_qualified_name)
Check if the given module is a leaf module.
Parameters:
- m (Module) –
- module_qualified_name (str) –
Return type:
bool
classmethod is_registered_leaf(m)
Check if module is registered as extra leaf module.
Parameters:
m (Module | type [ Module ]) –
Return type:
bool
is_unvisited(m)
Check if the node target is unvisited.
Parameters:
m (Module | Callable) –
Return type:
bool
record_call_module(m, e=None)
Record the call_module operation.
Parameters:
- m (Module) –
- e (Exception | None) –
Return type:
None
classmethod register_leaf(m)
Register a module as extra leaf module.
Parameters:
m (type [ Module ]) –
Return type:
None
trace(root, concrete_args=None)
Trace root using the fx tracer.
Unlike the standard tracer in fx, this tracer adds the following functionalities:
- Only torch.Tensors in concrete_args are converted to Proxy objects.
- The tracer doesn’t fail when encountering unsupported operations. Instead it will wrap the module containing the unsupported operation into a leaf node and repeat the tracing.
- The tracer doesn’t trace into leaf modules even if the leaf module is the root.
Parameters:
- root (Module) –
- concrete_args (dict [ str , Any ] | None) –
Return type:
Graph
classmethod unregister_leaf(m)
Unregister a module as extra leaf module.
Parameters:
m (type [ Module ]) –
Return type:
None
recursive_trace(model, concrete_args=None)
Recursively trace the model and all failing sub-modules and return graph collection.
Parameters:
- model (Module) – The model to be traced.
- concrete_args (dict [ str , Any ] | None) – [EXPERIMENTAL; not backwards compatible!] The concrete arguments to be used for tracing of the root module (i.e. the model’s forward method). Note that if/when the tracing recurses, no concrete args are passed on! Defaults to None.
Returns:
The collection of fx graphs resulting from the tracing.
Return type:
Note that the graph collection contains the graph of the root module as well as all failing submodules. For failing modules, the graph will contain the failing module as a leaf and then separate graphs for each submodule within the failing module to maximize the coverage of the tracing results.