def prepare_model_outputs(
name_a: str,
gm_a: GraphModule,
name_b: str,
gm_b: GraphModule,
logger_cls: Callable,
) -> Tuple[GraphModule, GraphModule]:
matched_subgraph_pairs = get_matching_subgraph_pairs(gm_a, gm_b)
nodes_and_names_to_instrument_a = []
nodes_and_names_to_instrument_b = []
for match_name, (subgraph_a, subgraph_b) in matched_subgraph_pairs.items():
node_start_a, node_end_a = subgraph_a
node_start_b, node_end_b = subgraph_b
# Note: for matching activations we always use the end nodes,
# such as observing the output of relu in linear-relu
nodes_and_names_to_instrument_a.append((node_end_a, match_name))
nodes_and_names_to_instrument_b.append((node_end_b, match_name))
gm_a = prepare_single_model_output(name_a, gm_a,
nodes_and_names_to_instrument_a,
logger_cls)
gm_b = prepare_single_model_output(name_b, gm_b,
nodes_and_names_to_instrument_b,
logger_cls)
return (gm_a, gm_b)
def test_simple_fun(self):
class M(nn.Module):
def __init__(self):
super().__init__()
self.w = nn.Parameter(torch.empty(1, 4))
self.b = nn.Parameter(torch.zeros(1))
torch.nn.init.kaiming_uniform_(self.w, a=math.sqrt(5))
def forward(self, x):
return F.linear(x, self.w, self.b)
m = M().eval()
mp = prepare_fx(m, {'': torch.quantization.default_qconfig})
# TODO(future PR): prevent the need for copying here, we can copy the
# modules but should reuse the underlying tensors
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
results = get_matching_subgraph_pairs(mp, mq)
expected_types = {
'base_op_torch.nn.functional.linear_0':
((F.linear, F.linear), (toq.linear, toq.linear))
}
self.assert_types_for_matched_subgraph_pairs(results, expected_types,
mp, mq)
def _extract_weights_impl(
model_name_a: str,
gm_a: GraphModule,
model_name_b: str,
gm_b: GraphModule,
base_name_to_sets_of_related_ops: Optional[Dict[
str, Set[NSNodeTargetType]]] = None,
unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None,
) -> NSResultsType:
matched_subgraph_pairs = get_matching_subgraph_pairs(
gm_a, gm_b, base_name_to_sets_of_related_ops, unmatchable_types_map)
# 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():
subgraph_a, subgraph_b = match
nodes_and_names_to_instrument_a.append(
(subgraph_a.base_op_node, match_name))
nodes_and_names_to_instrument_b.append(
(subgraph_b.base_op_node, match_name))
# populate the results, one model at a time
results: NSResultsType = {}
_extract_weights_one_model(model_name_a, gm_a,
nodes_and_names_to_instrument_a, results)
_extract_weights_one_model(model_name_b, gm_b,
nodes_and_names_to_instrument_b, results)
return results
def _add_shadow_loggers_impl(
name_a: str,
gm_a: GraphModule,
name_b: str,
gm_b: GraphModule,
logger_cls: Callable,
should_log_inputs: bool,
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:
matched_subgraph_pairs = get_matching_subgraph_pairs(
gm_a, gm_b, base_name_to_sets_of_related_ops, unmatchable_types_map)
gm_a_shadows_b = create_a_shadows_b(
name_a,
gm_a,
name_b,
gm_b,
matched_subgraph_pairs,
logger_cls,
should_log_inputs=should_log_inputs,
node_type_to_io_type_map=node_type_to_io_type_map)
return gm_a_shadows_b
def _extract_weights_impl(
model_name_a: str,
gm_a: GraphModule,
model_name_b: str,
gm_b: GraphModule,
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_impl")
matched_subgraph_pairs = get_matching_subgraph_pairs(
gm_a, gm_b, base_name_to_sets_of_related_ops,
unmatchable_types_map)
# 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():
subgraph_a, subgraph_b = match
nodes_and_names_to_instrument_a.append((subgraph_a.base_op_node, match_name))
nodes_and_names_to_instrument_b.append((subgraph_b.base_op_node, match_name))
# populate the results, one model at a time
results: NSResultsType = {}
_extract_weights_one_model(
model_name_a, gm_a, nodes_and_names_to_instrument_a, results)
_extract_weights_one_model(
model_name_b, gm_b, nodes_and_names_to_instrument_b, results)
# rekey on names of nodes in gm_b
results = rekey_logger_info_on_node_name_of_model(results, model_name_b)
return results
def _add_loggers_impl(
name_a: str,
gm_a: GraphModule,
name_b: str,
gm_b: GraphModule,
logger_cls: Callable,
should_log_inputs: bool,
) -> Tuple[nn.Module, nn.Module]:
matched_subgraph_pairs = get_matching_subgraph_pairs(gm_a, gm_b)
nodes_and_names_to_instrument_a = []
nodes_and_names_to_instrument_b = []
for match_name, (subgraph_a, subgraph_b) in matched_subgraph_pairs.items():
# Note: for matching activations we always use the end nodes,
# such as observing the output of relu in linear-relu
nodes_and_names_to_instrument_a.append(
(subgraph_a.end_node, match_name))
nodes_and_names_to_instrument_b.append(
(subgraph_b.end_node, match_name))
new_model_a = _add_loggers_one_model(name_a,
gm_a,
nodes_and_names_to_instrument_a,
logger_cls,
should_log_inputs=should_log_inputs)
new_model_b = _add_loggers_one_model(name_b,
gm_b,
nodes_and_names_to_instrument_b,
logger_cls,
should_log_inputs=should_log_inputs)
return (new_model_a, new_model_b)
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 test_nodes_before_cat(self):
# verify that nodes before cat get matched
class M(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x0):
x1 = torch.add(x0, 1.0)
y1 = torch.add(x0, 1.0)
x2 = torch.cat([x1, y1])
return x2
m = M().eval()
mp = prepare_fx(m, {'': torch.quantization.default_qconfig})
# TODO(future PR): prevent the need for copying here, we can copy the
# modules but should reuse the underlying tensors
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
results = get_matching_subgraph_pairs(mp, mq)
expected_types = {
'base_op_torch.cat_0':
((torch.cat, torch.cat), (toq.cat, toq.cat)),
'base_op_torch.add_0':
((torch.add, torch.add), (toq.add, toq.add)),
'base_op_torch.add_1':
((torch.add, torch.add), (toq.add, toq.add)),
}
self.assert_types_for_matched_subgraph_pairs(results, expected_types,
mp, mq)
def _add_loggers_impl(
name_a: str,
gm_a: GraphModule,
name_b: str,
gm_b: GraphModule,
logger_cls: Callable,
should_log_inputs: bool,
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]:
matched_subgraph_pairs = get_matching_subgraph_pairs(
gm_a, gm_b,
base_name_to_sets_of_related_ops, unmatchable_types_map)
nodes_and_names_to_instrument_inputs_a = []
nodes_and_names_to_instrument_inputs_b = []
nodes_and_names_to_instrument_outputs_a = []
nodes_and_names_to_instrument_outputs_b = []
for match_name, (subgraph_a, subgraph_b) in matched_subgraph_pairs.items():
# Note: for matching inputs we use start_node, such as observing
# the input of linear in linear-relu
if should_log_inputs:
nodes_and_names_to_instrument_inputs_a.append((subgraph_a.start_node, match_name))
nodes_and_names_to_instrument_inputs_b.append((subgraph_b.start_node, match_name))
# Note: for matching activations we always use end_node,
# such as observing the output of relu in linear-relu
nodes_and_names_to_instrument_outputs_a.append((subgraph_a.end_node, match_name))
nodes_and_names_to_instrument_outputs_b.append((subgraph_b.end_node, match_name))
new_model_a = _add_loggers_one_model(
name_a, gm_a, nodes_and_names_to_instrument_inputs_a,
nodes_and_names_to_instrument_outputs_a, logger_cls)
new_model_b = _add_loggers_one_model(
name_b, gm_b, nodes_and_names_to_instrument_inputs_b,
nodes_and_names_to_instrument_outputs_b, logger_cls)
return (new_model_a, new_model_b)
def test_dict_return_type(self):
# verify that we can traverse up nodes which return dictionaries
class M(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x0):
x1 = torch.add(x0, 1.0)
y1 = torch.add(x0, 1.0)
z1 = torch.add(x0, 1.0)
a1 = {'x1': x1, 'y1': (y1, ), 'z1': [{'key': (z1, )}]}
return a1
m = M().eval()
mp = prepare_fx(m, {'': torch.quantization.default_qconfig})
# TODO(future PR): prevent the need for copying here, we can copy the
# modules but should reuse the underlying tensors
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
results = get_matching_subgraph_pairs(mp, mq)
expected_types = {
'base_op_torch.add_0':
((torch.add, torch.add), (toq.add, toq.add)),
'base_op_torch.add_1':
((torch.add, torch.add), (toq.add, toq.add)),
'base_op_torch.add_2':
((torch.add, torch.add), (toq.add, toq.add)),
}
self.assert_types_for_matched_subgraph_pairs(results, expected_types,
mp, mq)
def test_matching_failure_node_type(self):
# verify that matching graphs with non-matching node types fails
m1 = nn.Sequential(nn.Conv2d(1, 1, 1)).eval()
m2 = nn.Sequential(nn.Linear(1, 1)).eval()
mp1 = prepare_fx(m1, {'': torch.quantization.default_qconfig})
mp2 = prepare_fx(m2, {'': torch.quantization.default_qconfig})
with self.assertRaises(GraphMatchingException) as ex:
results = get_matching_subgraph_pairs(mp1, mp2)
def test_mobilenet_v2_qat(self):
# verify that mobilenetv2 graph is able to be matched
import torchvision
m = torchvision.models.__dict__['mobilenet_v2'](pretrained=False).float()
mp = prepare_qat_fx(m, {'': torch.quantization.get_default_qat_qconfig('fbgemm')})
# TODO(future PR): prevent the need for copying here, we can copy the
# modules but should reuse the underlying tensors
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
# assume success if no exceptions
results = get_matching_subgraph_pairs(mp, mq)
def test_simple_mod(self):
m = nn.Sequential(nn.Conv2d(1, 1, 1)).eval()
mp = prepare_fx(m, {'': torch.quantization.default_qconfig})
# TODO(future PR): prevent the need for copying here, we can copy the
# modules but should reuse the underlying tensors
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
results = get_matching_subgraph_pairs(mp, mq)
expected_types = {'0': ((nn.Conv2d, nn.Conv2d), (nnq.Conv2d, nnq.Conv2d))}
self.assert_types_for_matched_subgraph_pairs(results, expected_types, mp, mq)
def test_simple_mod_multi(self):
m = nn.Sequential(
nn.Sequential(nn.Conv2d(1, 1, 1), ),
nn.Conv2d(1, 1, 1),
).eval()
mp = prepare_fx(m, {'': torch.quantization.default_qconfig})
# TODO(future PR): prevent the need for copying here, we can copy the
# modules but should reuse the underlying tensors
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
# assume success if no exceptions
results = get_matching_subgraph_pairs(mp, mq)
def _add_shadow_loggers_impl(
name_a: str,
gm_a: GraphModule,
name_b: str,
gm_b: GraphModule,
logger_cls: Callable,
should_log_inputs: bool,
) -> nn.Module:
matched_subgraph_pairs = get_matching_subgraph_pairs(gm_a, gm_b)
gm_a_shadows_b = create_a_shadows_b(
name_a, gm_a, name_b, gm_b, matched_subgraph_pairs, logger_cls,
should_log_inputs=should_log_inputs)
return gm_a_shadows_b
def test_simple_fusion(self):
m = LinearReluFunctional().eval()
mp = prepare_fx(m, {'': torch.quantization.default_qconfig})
# TODO(future PR): prevent the need for copying here, we can copy the
# modules but should reuse the underlying tensors
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
results = get_matching_subgraph_pairs(mp, mq)
expected_types = {
'base_op_torch.nn.functional.linear_0':
((F.linear, F.relu), (toq.linear_relu, toq.linear_relu)),
}
self.assert_types_for_matched_subgraph_pairs(results, expected_types, mp, mq)
def prepare_model_with_stubs(
name_a: str,
gm_a: GraphModule,
name_b: str,
gm_b: GraphModule,
logger_cls: Callable,
) -> GraphModule:
"""
Same thing as prepare_model_outputs, but for an `a_shadows_b` model.
TODO(future PR): real docblock
"""
matched_subgraph_pairs = get_matching_subgraph_pairs(gm_a, gm_b)
gm_a_shadows_b = create_a_shadows_b(name_a, gm_a, name_b, gm_b,
matched_subgraph_pairs, logger_cls)
return gm_a_shadows_b
def test_simple_tensor_ops(self):
class M(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x, y):
z = x + y
return z
m = M().eval()
mp = prepare_fx(m, {'': torch.quantization.default_qconfig})
# TODO(future PR): prevent the need for copying here, we can copy the
# modules but should reuse the underlying tensors
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
# assume success if no exceptions
results = get_matching_subgraph_pairs(mp, mq)
def test_nodes_with_equal_types_do_not_get_matched(self):
# verifies that by default, nodes with equivalent types do not get matched.
# This is important for user defined types, for which we do not know
# the weight extraction functions or input type. In the future, this can
# be made configurable.
class M(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 1, 1)
self.conv2 = nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = torch.mul(x, x)
x = torch.sigmoid(x)
x = F.relu(x)
return x
m = M().eval()
# prevent conv2 from getting quantized, so we can test
# modules with equal types
qconfig_dict = {
'': torch.quantization.default_qconfig,
'module_name': [('conv2', None)],
}
mp = prepare_fx(m, qconfig_dict)
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
results = get_matching_subgraph_pairs(mp, mq)
# Conv2 should not be matched because we disabled quantization for it,
# so its type is the same in mp and mq. sigmoid should not be
# matched because they use the same function in mp and mq. relu should
# be matched because it is in the allowlist of functions with same
# signature across dtypes.
expected_types = {
'base_op_torch.nn.Conv2d_0':
((nn.Conv2d, nn.Conv2d), (nnq.Conv2d, nnq.Conv2d)),
'base_op_torch.mul_0':
((torch.mul, torch.mul), (toq.mul, toq.mul)),
'base_op_torch.relu_0': ((F.relu, F.relu), (F.relu, F.relu)),
}
self.assert_types_for_matched_subgraph_pairs(results, expected_types,
mp, mq)
def prepare_model_outputs(
name_a: str,
gm_a: GraphModule,
name_b: str,
gm_b: GraphModule,
logger_cls: Callable,
) -> Tuple[GraphModule, GraphModule]:
matched_subgraph_pairs = get_matching_subgraph_pairs(gm_a, gm_b)
subgraphs_to_instrument_a = []
subgraphs_to_instrument_b = []
for match_name, (subgraph_a, subgraph_b) in matched_subgraph_pairs.items():
subgraphs_to_instrument_a.append((subgraph_a, match_name))
subgraphs_to_instrument_b.append((subgraph_b, match_name))
gm_a = prepare_single_model_output(name_a, gm_a, subgraphs_to_instrument_a,
logger_cls)
gm_b = prepare_single_model_output(name_b, gm_b, subgraphs_to_instrument_b,
logger_cls)
return (gm_a, gm_b)
def test_nodes_with_equal_types_get_matched(self):
class M(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 1, 1)
self.conv2 = nn.Conv2d(1, 1, 1)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = torch.mul(x, x)
x = torch.sigmoid(x)
x = F.relu(x)
return x
m = M().eval()
# prevent conv2 from getting quantized, so we can test
# modules with equal types
qconfig_dict = {
'': torch.quantization.default_qconfig,
'module_name': [('conv2', None)],
}
mp = prepare_fx(m, qconfig_dict)
mp_copy = copy.deepcopy(mp)
mq = convert_fx(mp_copy)
results = get_matching_subgraph_pairs(mp, mq)
# all of these should be matched
expected_types = {
'base_op_torch.nn.Conv2d_1':
((nn.Conv2d, nn.Conv2d), (nnq.Conv2d, nnq.Conv2d)),
'base_op_torch.nn.Conv2d_0':
((nn.Conv2d, nn.Conv2d), (nn.Conv2d, nn.Conv2d)),
'base_op_torch.mul_0':
((torch.mul, torch.mul), (toq.mul, toq.mul)),
'base_op_torch.relu_0': ((F.relu, F.relu), (F.relu, F.relu)),
'base_op_torch.sigmoid_0':
((torch.sigmoid, torch.sigmoid), (torch.sigmoid, torch.sigmoid)),
}
self.assert_types_for_matched_subgraph_pairs(results, expected_types,
mp, mq)
def _extract_weights_impl(
model_name_a: str,
gm_a: GraphModule,
model_name_b: str,
gm_b: GraphModule,
) -> NSResultsType:
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 = {}
_extract_weights_one_model(model_name_a, gm_a,
nodes_and_names_to_instrument_a, results)
_extract_weights_one_model(model_name_b, gm_b,
nodes_and_names_to_instrument_b, results)
return results