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]