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,
) -> NSResultsType:
torch._C._log_api_usage_once(
"quantization_api._numeric_suite_fx.extract_weights")
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))
gm_b = GraphModule(model_b, tracer_b.trace(model_b))
return _extract_weights_impl(model_name_a, gm_a, model_name_b, gm_b,
base_name_to_sets_of_related_ops,
unmatchable_types_map)
def compare_weights(
model_name_a: str,
gm_a: GraphModule,
model_name_b: str,
gm_b: GraphModule,
) -> NSResultsType:
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)
matched_subgraph_pairs = get_matching_subgraph_pairs(gm_a, gm_b)
# split the subgraph pairs into one data structure for each model
nodes_and_names_to_instrument_a: List[Tuple[Node, str]] = []
nodes_and_names_to_instrument_b: List[Tuple[Node, str]] = []
for match_name, match in matched_subgraph_pairs.items():
(node_start_a, node_end_a), (node_start_b, node_end_b) = match
nodes_and_names_to_instrument_a.append((node_start_a, match_name))
nodes_and_names_to_instrument_b.append((node_start_b, match_name))
# populate the results, one model at a time
results: NSResultsType = {}
add_weight_info_to_dict(model_name_a, gm_a,
nodes_and_names_to_instrument_a, results)
add_weight_info_to_dict(model_name_b, gm_b,
nodes_and_names_to_instrument_b, results)
return results
def compare_weights(
name_a: str,
gm_a: GraphModule,
name_b: str,
gm_b: GraphModule,
) -> Dict[str, Dict[str, torch.Tensor]]:
type_a_related_to_b = get_type_a_related_to_b()
matched_node_pairs = get_matching_node_pairs(gm_a, gm_b)
results = {}
for match_name, match in matched_node_pairs.items():
node_a, node_b = match
assert node_a.op == node_b.op and \
node_a.op in ('call_function', 'call_module')
if node_a.op == 'call_function':
# linear
# TODO(future PR): other function types
a_related_to_linear = node_a.target in (F.linear,) or \
(node_a.target, F.linear) in type_a_related_to_b
if a_related_to_linear:
weight_a = get_linear_fun_weight(node_a, gm_a)
weight_b = get_linear_fun_weight(node_b, gm_b)
results[match_name] = {
name_a: weight_a,
name_b: weight_b,
}
else: # call_module
# for call_module, we need to look up the modules to do the type check
assert isinstance(node_a.target, str)
mod_a = getattr_from_fqn(gm_a, node_a.target)
assert isinstance(node_b.target, str)
mod_b = getattr_from_fqn(gm_b, node_b.target)
# check that A is one the modules we need
# assume B is related (this is done by graph matcher)
a_related_to_conv2d_mod = isinstance(mod_a, nn.Conv2d) or \
(type(mod_a), nn.Conv2d) in type_a_related_to_b
# TODO(future PR): other module types
if a_related_to_conv2d_mod:
weight_a = get_conv_mod_weight(mod_a)
weight_b = get_conv_mod_weight(mod_b)
results[match_name] = {
name_a: weight_a,
name_b: weight_b,
}
return results
def add_weight_info_to_dict(
model_name: str,
model: GraphModule,
nodes_and_names_to_instrument: List[Tuple[Node, str]],
results: NSResultsType,
) -> 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)
for node, ref_name in nodes_and_names_to_instrument:
if ref_name not in results:
results[ref_name] = {}
if node.op == 'call_function':
# linear
# TODO(future PR): other function types
related_to_linear = node.target in (F.linear,) or \
(node.target, F.linear) in type_a_related_to_b
if related_to_linear:
weight = get_linear_fun_weight(node, model)
results[ref_name][model_name] = {
'type': NSSingleResultValuesType.WEIGHT.value,
'values': [weight],
'node_name': node.name,
'node_target_type': str(node.target),
}
else: # call_module
# for call_module, we need to look up the modules to do the type check
assert isinstance(node.target, str)
mod = getattr_from_fqn(model, node.target)
# check that A is one the modules we need
# assume B is related (this is done by graph matcher)
related_to_conv2d_mod = isinstance(mod, nn.Conv2d) or \
(type(mod), nn.Conv2d) in type_a_related_to_b
# TODO(future PR): other module types
if related_to_conv2d_mod:
weight = get_conv_mod_weight(mod)
results[ref_name][model_name] = {
'type': NSSingleResultValuesType.WEIGHT.value,
'values': [weight],
'node_name': node.name,
'node_target_type': str(type(mod)),
}
def extract_weights(
model_name_a: str,
model_a: nn.Module,
model_name_b: str,
model_b: nn.Module,
) -> NSResultsType:
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)
tracer_a, tracer_b = NSTracer(), NSTracer()
gm_a = GraphModule(model_a, tracer_a.trace(model_a))
gm_b = GraphModule(model_b, tracer_b.trace(model_b))
return _extract_weights_impl(model_name_a, gm_a, model_name_b, gm_b)
def _extract_weights_one_model(
model_name: str,
model: GraphModule,
nodes_and_names_to_instrument: List[Tuple[Node, str]],
results: NSResultsType,
) -> 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)
for node, ref_name in nodes_and_names_to_instrument:
res_type = NSSingleResultValuesType.WEIGHT.value
if ref_name not in results:
results[ref_name] = {res_type: {}}
extracted_weight = \
extract_weight_from_node(node, model, type_a_related_to_b)
if extracted_weight:
results[ref_name][res_type][model_name] = [extracted_weight]
def extract_weights(
model_name_a: str,
model_a: nn.Module,
model_name_b: str,
model_b: nn.Module,
) -> NSResultsType:
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))
gm_b = GraphModule(model_b, tracer_b.trace(model_b))
return _extract_weights_impl(model_name_a, gm_a, model_name_b, gm_b)
def _extract_weights_one_model(
model_name: str,
model: GraphModule,
nodes_and_names_to_instrument: List[Tuple[Node, str]],
results: NSResultsType,
) -> None:
torch._C._log_api_usage_once("quantization_api._numeric_suite_fx._extract_weights_one_model")
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)
for node, ref_name in nodes_and_names_to_instrument:
res_type = NSSingleResultValuesType.WEIGHT.value
extracted_weight = \
extract_weight_from_node(node, model, type_a_related_to_b)
if extracted_weight:
if ref_name not in results:
results[ref_name] = {res_type: {}}
results[ref_name][res_type][model_name] = [extracted_weight]
def add_weight_info_to_dict(
model_name: str,
model: GraphModule,
nodes_and_names_to_instrument: List[Tuple[Node, str]],
results: Dict[str, Dict[str, torch.Tensor]],
) -> None:
type_a_related_to_b = get_type_a_related_to_b()
for node, ref_node_name in nodes_and_names_to_instrument:
if ref_node_name not in results:
results[ref_node_name] = {}
if node.op == 'call_function':
# linear
# TODO(future PR): other function types
related_to_linear = node.target in (F.linear,) or \
(node.target, F.linear) in type_a_related_to_b
if related_to_linear:
weight = get_linear_fun_weight(node, model)
results[ref_node_name][model_name] = weight
else: # call_module
# for call_module, we need to look up the modules to do the type check
assert isinstance(node.target, str)
mod = getattr_from_fqn(model, node.target)
# check that A is one the modules we need
# assume B is related (this is done by graph matcher)
related_to_conv2d_mod = isinstance(mod, nn.Conv2d) or \
(type(mod), nn.Conv2d) in type_a_related_to_b
# TODO(future PR): other module types
if related_to_conv2d_mod:
weight = get_conv_mod_weight(mod)
results[ref_node_name][model_name] = weight
def _extract_weights_one_model(
model_name: str,
model: GraphModule,
nodes_and_names_to_instrument: List[Tuple[Node, str]],
results: NSResultsType,
) -> 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)
for node, ref_name in nodes_and_names_to_instrument:
res_type = NSSingleResultValuesType.WEIGHT.value
if ref_name not in results:
results[ref_name] = {res_type: {}}
if node.op == 'call_function':
# linear
# TODO(future PR): other function types
related_to_linear = node.target in (F.linear,) or \
(node.target, F.linear) in type_a_related_to_b
if related_to_linear:
weight = get_linear_fun_weight(node, model)
results[ref_name][res_type][model_name] = [{
'type':
res_type,
'values': [weight],
'prev_node_name':
node.name,
'prev_node_target_type':
str(node.target),
'ref_node_name':
node.name,
'index_within_arg':
0,
}]
else: # call_module
# for call_module, we need to look up the modules to do the type check
assert isinstance(node.target, str)
mod = getattr_from_fqn(model, node.target)
# check that A is one the modules we need
# assume B is related (this is done by graph matcher)
# TODO(future PR): 1d and 3d convs
related_to_conv2d_mod = isinstance(mod, nn.Conv2d) or \
(type(mod), nn.Conv2d) in type_a_related_to_b
related_to_linear_mod = isinstance(mod, nn.Linear) or \
(type(mod), nn.Linear) in type_a_related_to_b
related_to_lstm_mod = isinstance(mod, nn.LSTM) or \
(type(mod), nn.LSTM) in type_a_related_to_b
# TODO(future PR): other module types
if related_to_conv2d_mod:
weights = [get_conv_mod_weight(mod)]
elif related_to_lstm_mod:
weights = get_lstm_mod_weights(mod)
else:
assert related_to_linear_mod, f"module type {type(mod)} not handled yet"
weights = [get_linear_mod_weight(mod)]
results[ref_name][res_type][model_name] = [{
'type':
res_type,
'values':
weights,
'prev_node_name':
node.name,
'prev_node_target_type':
str(type(mod)),
'ref_node_name':
node.name,
'index_within_arg':
0,
}]
