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