Spaces:
Running
Running
import torch | |
from torch.fx._symbolic_trace import Tracer | |
from torch.fx.proxy import Scope | |
from torch.ao.nn.intrinsic import _FusedModule | |
from typing import List, Callable | |
__all__ = [ | |
"QuantizationTracer", | |
] | |
class ScopeContextManager(torch.fx.proxy.ScopeContextManager): | |
def __init__( | |
self, | |
scope: Scope, | |
current_module: torch.nn.Module, | |
current_module_path: str | |
): | |
super().__init__(scope, Scope(current_module_path, type(current_module))) | |
class QuantizationTracer(Tracer): | |
def __init__( | |
self, skipped_module_names: List[str], skipped_module_classes: List[Callable] | |
): | |
super().__init__() | |
self.skipped_module_names = skipped_module_names | |
self.skipped_module_classes = skipped_module_classes | |
# NB: initialized the module_type of top level module to None | |
# we are assuming people won't configure the model with the type of top level | |
# module here, since people can use "" for global config | |
# We can change this if there is a use case that configures | |
# qconfig using top level module type | |
self.scope = Scope("", None) | |
self.record_stack_traces = True | |
def is_leaf_module(self, m: torch.nn.Module, module_qualified_name: str) -> bool: | |
return ( | |
( | |
(m.__module__.startswith("torch.nn") or m.__module__.startswith("torch.ao.nn")) | |
and not isinstance(m, torch.nn.Sequential) | |
) | |
or module_qualified_name in self.skipped_module_names | |
or type(m) in self.skipped_module_classes | |
or isinstance(m, _FusedModule) | |
) | |