def add_loggers(
name_a: str,
model_a: nn.Module,
name_b: str,
model_b: nn.Module,
logger_cls: Callable,
should_log_inputs: bool = False,
base_name_to_sets_of_related_ops: Optional[Dict[
str, Set[NSNodeTargetType]]] = None,
unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None,
) -> Tuple[nn.Module, nn.Module]:
torch._C._log_api_usage_once(
"quantization_api._numeric_suite_fx.add_loggers")
# TODO(future PR): expose these
skipped_module_names: List[str] = []
skipped_module_classes: List[Callable] = []
tracer_a = NSTracer(skipped_module_names, skipped_module_classes)
tracer_b = NSTracer(skipped_module_names, skipped_module_classes)
gm_a = GraphModule(model_a, tracer_a.trace(model_a))
if hasattr(model_a, '_node_name_to_scope'):
gm_a._node_name_to_scope = model_a._node_name_to_scope
gm_b = GraphModule(model_b, tracer_b.trace(model_b))
if hasattr(model_b, '_node_name_to_scope'):
gm_b._node_name_to_scope = model_b._node_name_to_scope
return _add_loggers_impl(
name_a,
gm_a,
name_b,
gm_b,
logger_cls,
should_log_inputs=should_log_inputs,
base_name_to_sets_of_related_ops=base_name_to_sets_of_related_ops,
unmatchable_types_map=unmatchable_types_map)
def extract_weights(
model_name_a: str,
model_a: nn.Module,
model_name_b: str,
model_b: nn.Module,
base_name_to_sets_of_related_ops: Optional[Dict[
str, Set[NSNodeTargetType]]] = None,
unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None,
op_to_type_to_weight_extraction_fn: Optional[Dict[str,
Dict[Callable,
Callable]]] = None,
) -> NSResultsType:
torch._C._log_api_usage_once(
"quantization_api._numeric_suite_fx.extract_weights")
if base_name_to_sets_of_related_ops is None:
base_name_to_sets_of_related_ops = \
get_base_name_to_sets_of_related_ops()
type_a_related_to_b = \
get_type_a_related_to_b(base_name_to_sets_of_related_ops)
# TODO(future PR): expose these
skipped_module_names: List[str] = []
skipped_module_classes: List[Callable] = []
tracer_a = NSTracer(skipped_module_names, skipped_module_classes)
tracer_b = NSTracer(skipped_module_names, skipped_module_classes)
gm_a = GraphModule(model_a, tracer_a.trace(model_a))
if hasattr(model_a, '_node_name_to_scope'):
gm_a._node_name_to_scope = model_a._node_name_to_scope
gm_b = GraphModule(model_b, tracer_b.trace(model_b))
if hasattr(model_b, '_node_name_to_scope'):
gm_b._node_name_to_scope = model_b._node_name_to_scope
return _extract_weights_impl(model_name_a, gm_a, model_name_b, gm_b,
base_name_to_sets_of_related_ops,
unmatchable_types_map,
op_to_type_to_weight_extraction_fn)
def extract_weights(
model_name_a: str,
model_a: nn.Module,
model_name_b: str,
model_b: nn.Module,
base_name_to_sets_of_related_ops: Optional[Dict[
str, Set[NSNodeTargetType]]] = None,
unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None,
op_to_type_to_weight_extraction_fn: Optional[Dict[str,
Dict[Callable,
Callable]]] = None,
) -> NSResultsType:
"""
Extract weights from model A and model B, and return a comparison.
Args:
model_name_a: string name of model A to use in results
model_a: model A
model_name_b: string name of model B to use in results
model_b: model B
base_name_to_sets_of_related_ops: optional override of subgraph base nodes, subject to change
unmatchable_types_map: optional override of unmatchable types, subject to change
op_to_type_to_weight_extraction_fn: optional override of function which extracts weight
from a type, subject to change
Return:
NSResultsType, containing the weight comparisons
"""
torch._C._log_api_usage_once(
"quantization_api._numeric_suite_fx.extract_weights")
if base_name_to_sets_of_related_ops is None:
base_name_to_sets_of_related_ops = \
get_base_name_to_sets_of_related_ops()
type_a_related_to_b = \
get_type_a_related_to_b(base_name_to_sets_of_related_ops)
# TODO(future PR): expose these
skipped_module_names: List[str] = []
skipped_module_classes: List[Callable] = []
tracer_a = NSTracer(skipped_module_names, skipped_module_classes)
tracer_b = NSTracer(skipped_module_names, skipped_module_classes)
gm_a = GraphModule(model_a, tracer_a.trace(model_a))
if hasattr(model_a, '_node_name_to_scope'):
gm_a._node_name_to_scope = model_a._node_name_to_scope
gm_b = GraphModule(model_b, tracer_b.trace(model_b))
if hasattr(model_b, '_node_name_to_scope'):
gm_b._node_name_to_scope = model_b._node_name_to_scope
return _extract_weights_impl(model_name_a, gm_a, model_name_b, gm_b,
base_name_to_sets_of_related_ops,
unmatchable_types_map,
op_to_type_to_weight_extraction_fn)
def add_shadow_loggers(
name_a: str,
model_a: nn.Module,
name_b: str,
model_b: nn.Module,
logger_cls: Callable,
should_log_inputs: bool = False,
base_name_to_sets_of_related_ops: Optional[Dict[
str, Set[NSNodeTargetType]]] = None,
node_type_to_io_type_map: Optional[Dict[str,
Set[NSNodeTargetType]]] = None,
unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None,
) -> nn.Module:
"""
Instrument model A and model B with shadow loggers.
Args:
model_name_a: string name of model A to use in results
model_a: model A
model_name_b: string name of model B to use in results
model_b: model B
logger_cls: class of Logger to use
should_log_inputs: whether to log inputs
base_name_to_sets_of_related_ops: optional override of subgraph base nodes, subject to change
unmatchable_types_map: optional override of unmatchable types, subject to change
"""
torch._C._log_api_usage_once(
"quantization_api._numeric_suite_fx.add_shadow_loggers")
# TODO(future PR): expose these
skipped_module_names: List[str] = []
skipped_module_classes: List[Callable] = []
tracer_a = NSTracer(skipped_module_names, skipped_module_classes)
tracer_b = NSTracer(skipped_module_names, skipped_module_classes)
gm_a = GraphModule(model_a, tracer_a.trace(model_a))
if hasattr(model_a, '_node_name_to_scope'):
gm_a._node_name_to_scope = model_a._node_name_to_scope
gm_b = GraphModule(model_b, tracer_b.trace(model_b))
if hasattr(model_b, '_node_name_to_scope'):
gm_b._node_name_to_scope = model_b._node_name_to_scope
return _add_shadow_loggers_impl(
name_a,
gm_a,
name_b,
gm_b,
logger_cls,
should_log_inputs=should_log_inputs,
base_name_to_sets_of_related_ops=base_name_to_sets_of_related_ops,
node_type_to_io_type_map=node_type_to_io_type_map,
unmatchable_types_map=unmatchable_types_map)
def save_state(observed: GraphModule, qconfig_map: Dict[str, QConfigAny],
node_name_to_scope: Dict[str, Tuple[str, type]],
patterns: Dict[Pattern, QuantizeHandler],
prepare_custom_config_dict: Dict[str, Any]) -> None:
observed._patterns = patterns # type: ignore[assignment]
observed._qconfig_map = qconfig_map # type: ignore[assignment]
observed._prepare_custom_config_dict = \
prepare_custom_config_dict # type: ignore[assignment]
observed._node_name_to_scope = node_name_to_scope # type: ignore[assignment]
def save_state(
observed: GraphModule,
qconfig_map: Dict[str, QConfigAny],
node_name_to_scope: Dict[str, Tuple[str, type]],
patterns: Dict[Pattern, QuantizeHandler],
prepare_custom_config_dict: Dict[str, Any],
equalization_qconfig_map: Dict[str, Any],
qconfig_dict: Dict[str, Dict[Any, Any]],
is_training: bool,
) -> None:
observed._patterns = patterns # type: ignore[assignment]
observed._qconfig_map = qconfig_map # type: ignore[assignment]
observed._prepare_custom_config_dict = \
prepare_custom_config_dict # type: ignore[assignment]
observed._node_name_to_scope = node_name_to_scope # type: ignore[assignment]
observed._equalization_qconfig_map = equalization_qconfig_map # type: ignore[assignment]
observed._qconfig_dict = qconfig_dict # type: ignore[assignment]
observed._is_training = is_training # type: ignore[assignment]
