def test_graph_fns(self): g = Graph() a = g.placeholder('a') b = g.call_module('linear', (a, )) c = g.get_attr('bias') d = g.call_method('add', (b, c)) e = g.call_function(torch.sin, (d, )) g.output(e) mod = torch.nn.Module() mod.linear = torch.nn.Linear(3, 4) mod.bias = torch.rand(4) gm = GraphModule(mod, g) gm.graph.lint(gm) input = torch.rand(3) r = gm(input) ref = torch.sin(mod.linear(input) + mod.bias) self.assertEqual(r, ref)
def test_graph_unique_names(self): class M(torch.nn.Module): def forward(self, a, b): return a + b m = M() g = symbolic_trace(m).graph new_g = torch.fx.Graph() val_map : Dict[Node, Node] = {} output_val = new_g.graph_copy(g, val_map) t = Proxy(output_val) # test that we can use proxy objects to generate more graph code later for things that do not need to work with modules. new_g.output((t + t).node) gm = GraphModule(m, new_g) seen_names : Set[str] = set() for node in gm.graph.nodes: assert node.name not in seen_names seen_names.add(node.name)
def _prepare_fx(model, qconfig_dict, inplace, prepare_custom_config_dict=None, is_standalone_module=False): r""" Internal helper function for prepare_fx Args: `model`, `qconfig_dict`, `inplace` `prepare_custom_config_dict`: see docs for :func:`~torch.quantization.prepare_fx` `is_standalone_module`: a boolean flag indicates whether we are quantizing a standalone module or not, a standalone module is a submodule of the parent module that is not inlined in the forward graph of the parent module, the way we quantize standalone module is described in: :func:`~torch.quantization._prepare_standalone_module_fx` """ if prepare_custom_config_dict is None: prepare_custom_config_dict = {} skipped_module_names = prepare_custom_config_dict.get( "non_traceable_module_name", []) skipped_module_classes = prepare_custom_config_dict.get( "non_traceable_module_class", []) # swap FloatFunctional with FXFloatFunctional _swap_ff_with_fxff(model) # symbolically trace the model if not is_standalone_module: # standalone module and custom module config are applied in top level module standalone_module_names = prepare_custom_config_dict.get( 'standalone_module_name', []) skipped_module_names += standalone_module_names custom_module_config = prepare_custom_config_dict.get( 'float_to_observed_custom_module_class', {}) custom_module_classes = list(custom_module_config.keys()) skipped_module_classes += custom_module_classes tracer = CustomTracer(skipped_module_names, skipped_module_classes) graph_module = GraphModule(model, tracer.trace(model)) graph_module = _fuse_fx(graph_module, inplace) quantizer = Quantizer() return quantizer.prepare( graph_module, qconfig_dict, inplace=True, prepare_custom_config_dict=prepare_custom_config_dict, is_standalone_module=is_standalone_module)
def quantize(self): self.quantized_graph = Graph() env = {} quant_env = {} def load_arg(n, quantized): if not quantized: if n.name not in env and n.name in quant_env: env[n.name] = Proxy(quant_env[n.name]).dequantize().node return env[n.name] else: if n.name not in quant_env and n.name in env: quant_env[n.name] = self.quants[n.name].quantize( env[n.name]) return quant_env[n.name] def copy_recursive(node): def load_or_emit(n): if n.name in env or e.name in quant_env: return load_arg(n, quantized=False) else: return copy_recusive(n) r = env[node.name] = self.quantized_graph.node_copy( node, lambda n: load_arg(n, quantized=False)) return r for node in self.graph.nodes: root_node, obj = self.matches.get(node.name, (None, None)) if root_node is None: # not quantized just copy it env[node.name] = self.quantized_graph.node_copy( node, lambda n: load_arg(n, quantized=False)) elif root_node is node: r = obj.quantize( self, node, lambda a: map_arg( a, lambda n: load_arg(n, quantized=True))) if r is NotImplemented: # quantizer choose to to quantize the node take the entire match, and just copy it over env[node.name] = copy_recursive(node) else: quant_env[node.name] = r return GraphModule(self.root, self.quantized_graph)
def _fold_weight(self, quantized): packed_weights = dict() # map from folded node name to the prepacked weight name folded_nodes = dict() # get packed weights for node in quantized.graph.nodes: if node.op == 'call_function' and node.target in WEIGHT_PREPACK_OPS: nodes_to_fold = collect_producer_nodes(node) if nodes_to_fold is not None: for node_to_fold in nodes_to_fold: folded_nodes[node_to_fold.name] = node prepacking_module = graph_module_from_producer_nodes( quantized, nodes_to_fold) packed_weight = prepacking_module() packed_weights[node.name] = packed_weight # remove folded nodes and replace the prepacking node with getattr folded_graph = Graph() env = {} def load_arg(a): return map_arg(a, lambda node: env[node.name]) get_new_packed_weight_name = get_new_attr_name_with_prefix( '_fx_pass_packed_weight_') quantized_root = quantized quantized_graph = quantized.graph for node in quantized_graph.nodes: prepack_node = folded_nodes.get(node.name, None) if prepack_node is node: packed_weight = packed_weights[node.name] # add a prepacked attribute to root packed_weight_name = get_new_packed_weight_name(quantized_root) setattr(quantized_root, packed_weight_name, packed_weight) # replace prepack node with a getattr node env[node.name] = folded_graph.create_node( 'get_attr', packed_weight_name, (), {}) elif prepack_node is not None: # remove the foled node continue else: # copy other nodes env[node.name] = folded_graph.node_copy(node, load_arg) quantized = GraphModule(quantized_root, folded_graph) return quantized
def _transform_remove_duplicates(module: GraphModule, debug: bool) -> GraphModule: """Removes duplicate modules by creating a copy of the module. This is necessary because BackPACK saves input/output which is overwritten if the module is called multiple times. Args: module: container module to transform debug: whether to print debug messages Returns: equivalent transformed module Raises: NotImplementedError: if a duplicate module has parameters """ if debug: print("\tBegin transformation: remove duplicates") graph: Graph = BackpackTracer().trace(module) targets = [n.target for n in graph.nodes] duplicates = {t for t in targets if targets.count(t) > 1} nodes = [n for n in graph.nodes if n.target in duplicates] for node in nodes: target = node.target original_module = module.get_submodule(target) for _ in original_module.parameters(): raise NotImplementedError( f"Cycle with parameters detected: module {original_module} with target" f" {target} has parameters and is used {targets.count(target)} times." ) new_module = deepcopy(original_module) new_target = _get_free_name(module, target) module.add_submodule(new_target, new_module) node.target = new_target graph.lint() if debug: print(f"\tDuplicates removed: {len(nodes)}") return GraphModule(module, graph)
def test_type_check_conv2D_maxpool2d_flatten(self): class BasicBlock(torch.nn.Module): def __init__(self): super(BasicBlock, self).__init__() self.conv1 = torch.nn.Conv2d(3, 6, 5) self.pool = torch.nn.MaxPool2d(2, 2) self.conv2 = torch.nn.Conv2d(6, 16, 5) self.fc1 = torch.nn.Linear(5, 120) self.pool2 = torch.nn.AdaptiveAvgPool2d((6, 7)) def forward(self, x: TensorType((4, 3, 32, 32))): out = self.conv1(x) out = self.pool(out) out = self.conv2(out) out = self.pool(out) out = self.fc1(out) out = self.pool2(out) out = torch.flatten(out, 1) return out B = BasicBlock() ast_rewriter = RewritingTracer() graph = ast_rewriter.trace(B) traced = GraphModule(ast_rewriter.root, graph, "gm") tc = GraphTypeChecker({}, traced) tc.type_check() expected_ph_types = [ TensorType((4, 3, 32, 32)), TensorType((4, 6, 28, 28)), TensorType((4, 6, 14, 14)), TensorType((4, 16, 10, 10)), TensorType((4, 16, 5, 5)), TensorType((4, 16, 5, 120)), TensorType((4, 16, 6, 7)), TensorType((4, 672)), TensorType((4, 672)) ] expected_iter = iter(expected_ph_types) traced.graph.eliminate_dead_code() for n in traced.graph.nodes: assert n.type == next(expected_iter)
def remove_observers_add_loggers( gm: GraphModule, node_to_instrument_to_ref_node_name: Dict[Node, Optional[str]], logger_cls: Callable, model_name: str, ) -> GraphModule: """ Takes the graph of gm, removes all observers, adds loggers to the output of each node in nodes_to_instrument. Returns a GraphModule with the new graph. """ new_graph = Graph() env: Dict[str, Any] = {} modules = dict(gm.named_modules()) def load_arg(a): return map_arg(a, lambda node: env[node.name]) for node in gm.graph.nodes: if node.op == 'output': new_graph.output(map_arg(node.args[0], load_arg)) continue if node.op == 'call_module' and is_activation_post_process( modules[node.target]): # remove activation post process node env[node.name] = env[node.args[0].name] elif node in node_to_instrument_to_ref_node_name: other_node_name = node_to_instrument_to_ref_node_name[node] # ensure env is populated with base node env[node.name] = new_graph.node_copy(node, load_arg) # add the logger after the base node env[node.name] = _insert_logger_after_node(env[node.name], gm, logger_cls, '_ns_logger_', model_name, other_node_name) else: env[node.name] = new_graph.node_copy(node, load_arg) new_gm = GraphModule(gm, new_graph) return new_gm
def _prepare_fx(model: torch.nn.Module, qconfig_dict: Any, prepare_custom_config_dict: Dict[str, Any] = None, is_standalone_module: bool = False) -> GraphModule: r""" Internal helper function for prepare_fx Args: `model`, `qconfig_dict`, `prepare_custom_config_dict`: see docs for :func:`~torch.quantization.prepare_fx` `is_standalone_module`: a boolean flag indicates whether we are quantizing a standalone module or not, a standalone module is a submodule of the parent module that is not inlined in the forward graph of the parent module, the way we quantize standalone module is described in: :func:`~torch.quantization._prepare_standalone_module_fx` """ if prepare_custom_config_dict is None: prepare_custom_config_dict = {} skipped_module_names = prepare_custom_config_dict.get("non_traceable_module_name", []) skipped_module_classes = prepare_custom_config_dict.get("non_traceable_module_class", []) # swap FloatFunctional with FXFloatFunctional _swap_ff_with_fxff(model) # symbolically trace the model if not is_standalone_module: # standalone module and custom module config are applied in top level module standalone_module_name_configs = prepare_custom_config_dict.get("standalone_module_name", []) skipped_module_names += [config[0] for config in standalone_module_name_configs] standalone_module_class_configs = prepare_custom_config_dict.get("standalone_module_class", []) skipped_module_classes += [config[0] for config in standalone_module_class_configs] float_custom_module_classes = get_custom_module_class_keys( prepare_custom_config_dict, "float_to_observed_custom_module_class") skipped_module_classes += float_custom_module_classes tracer = CustomTracer(skipped_module_names, skipped_module_classes) graph_module = GraphModule(model, tracer.trace(model)) graph_module = _fuse_fx(graph_module, prepare_custom_config_dict) quantizer = Quantizer() return quantizer.prepare( graph_module, qconfig_dict, prepare_custom_config_dict=prepare_custom_config_dict, is_standalone_module=is_standalone_module)
def test_type_typechecl_maxpool2d_3dinput(self): class BasicBlock(torch.nn.Module): def __init__(self): super(BasicBlock, self).__init__() self.pool = torch.nn.MaxPool2d(5, 8) def forward(self, x: TensorType((64, 8, 8))): out = self.pool(x) return out B = BasicBlock() ast_rewriter = RewritingTracer() graph = ast_rewriter.trace(B) traced = GraphModule(ast_rewriter.root, graph, "gm") tc = GraphTypeChecker({}, traced) tc.type_check() for n in traced.graph.nodes: if n.target == 'output': assert n.type == TensorType((64, 1, 1))
def test_type_check_batch_norm_2D_false(self): class BasicBlock(torch.nn.Module): def __init__(self, inplanes, planes): super(BasicBlock, self).__init__() norm_layer = torch.nn.BatchNorm2d self.bn1 = norm_layer(planes) def forward(self, x: TensorType((2, 2, 5))): identity = x out: TensorType((2, 2, Dyn, 4)) = self.bn1(x) out += identity return out B = BasicBlock(2, 2) ast_rewriter = RewritingTracer() graph = ast_rewriter.trace(B) traced = GraphModule(ast_rewriter.root, graph, "gm") tc = GraphTypeChecker({}, traced) with self.assertRaises(TypeError): tc.type_check()
def _prepare_fx(model, qconfig_dict, inplace, prepare_custom_config_dict=None, is_standalone_module=False): r""" Internal helper function for prepare_fx Args: `model`, `qconfig_dict`, `inplace` `prepare_custom_config_dict`: see docs for :func:`~torch.quantization.prepare_fx` `is_standalone_module`: a boolean flag indicates whether we are quantizing a standalone module or not, a standalone module is a submodule of the parent module that is not inlined in the forward graph of the parent module, the way we quantize standalone module is described in: :func:`~torch.quantization._prepare_standalone_module_fx` """ if prepare_custom_config_dict is None: prepare_custom_config_dict = {} # symbolically trace the model if is_standalone_module: # standlone module is traced before quantizing standalone modules graph_module = symbolic_trace(model) else: standalone_modules = prepare_custom_config_dict.get( 'standalone_module_name', []) custom_module_config = qconfig_dict.get('custom_module_class', []) custom_module_classes = [config[0] for config in custom_module_config] # TODO: currently we are registering classes globally, # we want to make custom module class mapping local _register_custom_module_class(custom_module_config) # skipping tracing standalone modules when tracing top level module tracer = CustomTracer(standalone_modules, custom_module_classes) graph_module = GraphModule(model, tracer.trace(model)) graph_module = _fuse_fx(graph_module, inplace) quantizer = Quantizer() return quantizer.prepare( graph_module, qconfig_dict, inplace=True, prepare_custom_config_dict=prepare_custom_config_dict, is_standalone_module=is_standalone_module)
def _transform_lstm_rnn(module: Module, debug: bool) -> GraphModule: """Transforms multi-layer RNN/LSTM to Sequential of single-layer RNN/LSTM. Converts multi-layer RNN/LSTM to Sequential with single-layer RNN/LSTM. If dropout probability is nonzero, creates intermediate dropout layers. Finally, copies training mode. Args: module: container module to transform debug: whether to print debug messages Returns: equivalent transformed module Raises: NotImplementedError: if initial hidden state is used in forward pass """ if debug: print("\tBegin transformation: LSTM, RNN") graph: Graph = BackpackTracer().trace(module) nodes = [ n for n in graph.nodes if n.op == "call_module" and isinstance(module.get_submodule(n.target), (RNN, LSTM)) and module.get_submodule(n.target).num_layers > 1 ] for node in nodes: if len(node.args) > 1: raise NotImplementedError( "For conversion, LSTM/RNN input must not have hidden states." ) lstm_module_replace = _make_rnn_backpack(module.get_submodule(node.target)) module.add_module(node.target, lstm_module_replace) graph.lint() if debug: print(f"\tRNNs, LSTMs transformed: {len(nodes)}") return GraphModule(module, graph)
def graph_module_from_producer_nodes(root, producer_nodes): r''' Construct a graph module from extracted producer nodes from `collect_producer_nodes` function Args: root: the root module for the original graph producer_nodes: a list of nodes we use to construct the graph Return: A graph module constructed from the producer nodes ''' assert len(producer_nodes) > 0, 'list of producer nodes can not be empty' # since we traced back from node to getattrr producer_nodes.reverse() graph = Graph() env = {} def load_arg(a): return map_arg(a, lambda node: env[node.name]) for producer_node in producer_nodes: env[producer_node.name] = graph.node_copy(producer_node, load_arg) graph.output(load_arg(producer_nodes[-1].name)) graph_module = GraphModule(root, graph) return graph_module
def _transform_transpose_to_module(module: Module, debug: bool) -> GraphModule: """Transforms transpose function or method to Permute module. The Permute module is initialized with transpose parameters and computes the permutation on its first forward pass. Args: module: container module to transform debug: whether to print debug messages Returns: equivalent transformed module """ target_function = "<built-in method transpose" target_method = "transpose" if debug: print(f"\tBegin transformation: {target_method} -> Permute") graph: Graph = BackpackTracer().trace(module) nodes = [ n for n in graph.nodes if (n.op == "call_function" and target_function in str(n.target)) or (n.op == "call_method" and target_method == str(n.target)) ] for node in nodes: _change_node_to_module( node, "permute", module, Permute(*node.args[1:], init_transpose=True), (node.args[0],), ) graph.lint() if debug: print(f"\tPermute transformed: {len(nodes)}") return GraphModule(module, graph)
def _transform_flatten_to_module(module: Module, debug: bool) -> GraphModule: """Transforms PyTorch's flatten method to the nn.Flatten module. Args: module: container module to transform debug: whether to print debug messages Returns: equivalent transformed module """ target_function = "<built-in method flatten" target_method = "flatten" if debug: print(f"\tBegin transformation: {target_function} -> Flatten") graph: Graph = BackpackTracer().trace(module) nodes = [ n for n in graph.nodes if (n.op == "call_function" and target_function in str(n.target)) or (n.op == "call_method" and target_method == str(n.target)) ] for node in nodes: start_dim = node.args[1] if len(node.args) > 1 else 0 end_dim = node.args[2] if len(node.args) > 2 else -1 _change_node_to_module( node, "flatten", module, Flatten(start_dim, end_dim), (node.args[0],) ) graph.lint() if debug: print(f"\tFlatten functions transformed: {len(nodes)}") return GraphModule(module, graph)
def _transform_permute_to_module(module: Module, debug: bool) -> GraphModule: """Transforms permute function or method to Permute module. Args: module: container module to transform debug: whether to print debug messages Returns: equivalent transformed module """ target1 = "permute" target2 = "<built-in method permute" if debug: print(f"\tBegin transformation: {target1}|{target2} -> Permute") graph: Graph = BackpackTracer().trace(module) nodes = [ n for n in graph.nodes if (n.op == "call_function" and target2 in str(n.target)) or (n.op == "call_method" and target1 == str(n.target)) ] for node in nodes: _change_node_to_module( node, "permute", module, Permute(*node.args[1]) if len(node.args) == 2 else Permute(*node.args[1:]), (node.args[0],), ) graph.lint() if debug: print(f"\tPermute transformed: {len(nodes)}") return GraphModule(module, graph)
def test_type_check_conv2D_2_fully_static(self): annotation_list = [(1, 2, 3, 5), (2, 5, 6, 9), (10, 15, 13, 14), (10, Dyn, 13, 14), (Dyn, Dyn, Dyn, 3)] input_list = [(1, 2, 3, 5), (2, 5, 6, 9), (10, 15, 13, 14), (10, 15, 13, 14), (1, 2, 2, 3)] intermediate_types = [(1, Dyn, Dyn, 7), (2, Dyn, 4, 6), (10, 15, Dyn, 5), (10, 15, 7, 7), (1, Dyn, Dyn, Dyn)] in_planes_list = [2, 5, 15, 15, 2] stride_list = [1, 2, 3, 2, 2] out_planes_list = [2, 5, 15, 15, 2] groups_list = [1, 5, 5, 5, 2] dilation_list = [1, 2, 3, 3, 3] padding_list = [1, 2, 3, 3, 3] kernel_size_list = [1, 2, 3, 3, 3] output_types = [(1, 2, Dyn, 7), (2, 5, 4, 6), (10, 15, Dyn, 5), (10, 15, 7, 7), (1, 2, Dyn, Dyn)] for i in range(5): annotation = annotation_list[i] input = input_list[i] in_planes = in_planes_list[i] stride = stride_list[i] out_planes = out_planes_list[i] groups = groups_list[i] dilation = dilation_list[i] padding = padding_list[i] kernel_size = kernel_size_list[i] intermediate_type = intermediate_types[i] class BasicBlock(torch.nn.Module): def __init__(self, in_planes, out_planes, kernel_size, stride, padding, groups, dilation): super(BasicBlock, self).__init__() self.conv1 = torch.nn.Conv2d(in_channels=in_planes, out_channels=out_planes, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=False, dilation=dilation) def forward(self, x): out = self.conv1(x) return out B = BasicBlock(in_planes, out_planes, kernel_size, stride, padding, groups, dilation) ast_rewriter = RewritingTracer() graph = ast_rewriter.trace(B) traced = GraphModule(ast_rewriter.root, graph, "gm") # annotate our argument for n in graph.nodes: if n.op == 'placeholder': n.type = TensorType(annotation) b = B.forward(torch.rand(input)) tc = GraphTypeChecker({}, traced) tc.type_check() for n in graph.nodes: if n.op == 'output': assert is_consistent(n.type, TensorType(b.size())) # test with intermediate annotations class BasicBlock(torch.nn.Module): def __init__(self, in_planes, out_planes, kernel_size, stride, padding, groups, dilation): super(BasicBlock, self).__init__() self.conv1 = torch.nn.Conv2d(in_channels=in_planes, out_channels=out_planes, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=False, dilation=dilation) def forward(self, x): out = self.conv1(x) return out B = BasicBlock(in_planes, out_planes, kernel_size, stride, padding, groups, dilation) ast_rewriter = RewritingTracer() graph = ast_rewriter.trace(B) traced = GraphModule(ast_rewriter.root, graph, "gm") # populate our intermediate notes for n in traced.graph.nodes: if n.op == 'call_module': n.type = TensorType(intermediate_type) tc = GraphTypeChecker({}, traced) tc.type_check() for n in traced.graph.nodes: if n.op == 'output': assert n.type == TensorType(output_types[i]) assert is_consistent(n.type, TensorType(b.size()))
def _convert(self, model, inplace=False, debug=False, is_dynamic_quant=False): self.restore_state(model) if not inplace: model = copy.deepcopy(model) self.is_dynamic_quant = is_dynamic_quant # run weight observers before inserting quant dequant nodes # for dynamic quantization if self.is_dynamic_quant: self._run_weight_observers(model) # move to cpu since we only have quantized cpu kernels model.eval().cpu() self.modules = dict(model.named_modules()) matches = self._find_matches(model.graph, self.modules, self.patterns) quants = self._find_quants(model.graph, matches) self.quantized_graph = Graph() env = {} quant_env = {} def load_non_quantized(n): if n.name not in env: assert n.name in quant_env, \ 'trying to load float node but did not find node:' + n.name + \ ' in quantized or non quantized environment, env: ' + str(env) + \ ' quant_env:' + str(quant_env) env[n.name] = Proxy(quant_env[n.name]).dequantize().node return env[n.name] def load_quantized(n): if n.name not in quant_env: assert n.name in env, \ 'trying to load quantized node but did not find node:' + n.name + \ ' in float environment:' + str(env) assert n.name in quants, 'did not find quant object for node:' + n.name quant = quants[n.name][0] quant_env[n.name] = quant.convert(self, env[n.name]) return quant_env[n.name] def load_x(n): assert n.name in env or n.name in quant_env, \ 'node ' + n.name + ' does not exist in either environment' if n.name in quant_env: return quant_env[n.name] else: return env[n.name] def load_arg(quantized): """ Input: quantized, which can be None, list, boolean or tuple - if quantized is a list or tuple, then arg should be a list and the args with corresponding indexes will be quantized - if quantized is a boolean, then all args will be quantized/not quantized - if quantized is None, then we'll load the node as long as it exists Output: fn which takes arg_or_args, and loads them from the corresponding environment depending on the value of quantized. """ assert quantized is None or isinstance(quantized, (tuple, list, bool)), type(quantized) def load_arg_impl(arg_or_args): if quantized is None: return map_arg(arg_or_args, load_x) if isinstance(quantized, bool): return map_arg(arg_or_args, load_quantized if quantized else load_non_quantized) elif isinstance(quantized, (tuple, list)): assert isinstance(arg_or_args, (tuple, list)), arg_or_args loaded_args = [] # for now, we only support quantizing positional arguments for i, a in enumerate(arg_or_args): if i in quantized: loaded_args.append(map_arg(a, load_quantized)) else: loaded_args.append(map_arg(a, load_non_quantized)) return type(arg_or_args)(loaded_args) return load_arg_impl def is_quantized(node): if isinstance(node, Node): assert node.name in env or node.name in quant_env, 'Expecting node to be in the environment' # there might be nodes appearing in both environemnts, but quant_env will take # precedence if node.name in quant_env: return True elif node.name in env: return False elif isinstance(node, list): quantized = map(is_quantized, node) if all(quantized): return True elif not any(quantized): return False else: raise Exception("partially quantized inputs in list not handled yet") for node in model.graph.nodes: root_node, matched, obj, qconfig = matches.get(node.name, (None, None, None, None)) if root_node is node: if qconfig is None: result = self.quantized_graph.node_copy(node, load_non_quantized) quantized = False else: result = obj.convert(self, node, load_arg) # Need to get correct quantized/non-quantized state for the output of CopyNode if isinstance(obj, CopyNode): assert node.op in [ 'call_module', 'call_function', 'call_method'], \ 'CopyNode of type ' + node.op + ' is not handled' quantized = is_quantized(node.args[0]) else: quantized = True # output of dynamic quantization is not quantized if self.is_dynamic_quant: quantized = False if quantized: quant_env[node.name] = result else: env[node.name] = result continue elif root_node is not None: continue # handle activation post process calls if node.op == 'call_module': if node.target.split('.')[-1].startswith('activation_post_process_'): observer_module = self.modules[node.target] prev_node = node.args[0] if observer_module.dtype == torch.float16: # activations are not quantized for # fp16 dynamic quantization # copy the activaiton_post_process node here # since we may need it when we insert prepack # op for weight of linear, this will be removed # later in a separate pass env[node.name] = self.quantized_graph.node_copy(node, load_non_quantized) continue if prev_node.name in quant_env: # if previous node is already quantized, we'll just remove the activation_post_process quant_env[node.name] = quant_env[prev_node.name] continue # replace activation post process with quantization ops root_module = self.modules[''] quant_env[node.name] = quantize_node( root_module, self.quantized_graph, load_non_quantized(node.args[0]), observer_module) continue # dequantize inputs for the node that are not quantized env[node.name] = self.quantized_graph.node_copy(node, load_non_quantized) self.quantized_graph.output(map_arg(model.graph.result, load_non_quantized)) # remove activation post process act_post_process_removed_graph = Graph() env = {} def load_arg(a): return map_arg(a, lambda node: env[node.name]) for node in self.quantized_graph.nodes: if node.op == 'call_module' and \ node.target.split('.')[-1].startswith('activation_post_process_'): # remove activation post process env[node.name] = env[node.args[0].name] else: env[node.name] = act_post_process_removed_graph.node_copy(node, load_arg) act_post_process_removed_graph.output(map_arg(self.quantized_graph.result, load_arg)) to_be_removed = [] for name, _ in model.named_modules(): if name.split('.')[-1].startswith('activation_post_process_'): to_be_removed.append(name) for n in to_be_removed: delattr(model, n) _remove_qconfig(model) model = GraphModule(model, act_post_process_removed_graph) return model
def _prepare(self, model, qconfig_dict, inplace, is_dynamic_quant): if not inplace: model = copy.deepcopy(model) self.is_dynamic_quant = is_dynamic_quant # TODO: allow user specified patterns if self.is_dynamic_quant: self.patterns = get_dynamic_quant_patterns() else: self.patterns = get_quant_patterns() propagate_qconfig_(model, qconfig_dict) if model.training: self._qat_swap_modules(model) self.modules = dict(model.named_modules()) # map from node name to qconfig, used in _find_matches self._generate_qconfig_map(model, model.graph) # match the patterns that will get quantized matches = self._find_matches(model.graph, self.modules, self.patterns) # find _inputs_ to matched nodes that are not quantized, these # have to be quantized, which requires measuring stats, # initialize an DefaultQuant object for each quants = self._find_quants(model.graph, matches) self.activation_post_process_map = dict() env = {} observed_graph = Graph() observed_node_names_set = set() def load_arg(a): return map_arg(a, lambda node: env[node.name]) for node in model.graph.nodes: if node.name in observed_node_names_set: continue get_new_observer_name = get_new_attr_name_with_prefix('activation_post_process_') root_node, _, obj, qconfig = matches.get(node.name, (None, None, None, None)) if root_node is None: env[node.name] = observed_graph.node_copy(node, load_arg) elif root_node is node: env[node.name] = observed_graph.node_copy(node, load_arg) def insert_observer(node, observer, device): observer_name = get_new_observer_name(model) setattr(model, observer_name, observer) self.activation_post_process_map[node.name] = observer env[node.name] = observed_graph.create_node('call_module', observer_name, (load_arg(node),), {}) observed_node_names_set.add(node.name) if device: getattr(model, observer_name).to(device) # don't need to insert observer for output in dynamic quantization if self.is_dynamic_quant: continue if isinstance(obj, CopyNode): assert node.op in [ 'call_module', 'call_function', 'call_method'], \ 'CopyNode of type ' + node.op + ' is not handled' def is_observed(input_arg): if isinstance(input_arg, Node): return input_arg.name in observed_node_names_set elif isinstance(input_arg, list): return all(map(is_observed, input_arg)) # propagate observed property from input if is_observed(node.args[0]): observed_node_names_set.add(node.name) elif (isinstance(obj, Add) or isinstance(obj, Mul)) and not obj.all_nodes: if node.args[0].name in observed_node_names_set: observed_node_names_set.add(node.name) elif qconfig is not None and obj.all_nodes: # observer for outputs new_observer = qconfig.activation() # respect device affinity when adding observers device = assert_and_get_unique_device(model) insert_observer(node, new_observer, device) else: env[node.name] = observed_graph.node_copy(node, load_arg) if node.name not in observed_node_names_set and node.name in quants: observer_name = get_new_observer_name(model) _, qconfig, is_weight = quants[node.name] if qconfig is not None: new_observer = \ qconfig.weight() if is_weight else qconfig.activation() # respect device affinity when adding observers device = assert_and_get_unique_device(model) if device: new_observer.to(device) self.activation_post_process_map[node.name] = new_observer setattr(model, observer_name, self.activation_post_process_map[node.name]) env[node.name] = observed_graph.create_node('call_module', observer_name, (load_arg(node),), {}) observed_node_names_set.add(node.name) observed_graph.output(load_arg(model.graph.result)) model = GraphModule(model, observed_graph) self.save_state(model) return model
def _prepare(self, model, qconfig_dict, inplace, quant_type): input_root = model.root if not inplace: input_root = copy.deepcopy(input_root) input_graph = model.graph self.quant_type = quant_type # TODO: allow user specified patterns if self.quant_type == QuantType.DYNAMIC: self.patterns = get_dynamic_quant_patterns() else: self.patterns = get_quant_patterns() propagate_qconfig_(input_root, qconfig_dict) if input_root.training: self._qat_swap_modules(input_root) self.modules = dict(input_root.named_modules()) # map from node name to qconfig, used in _find_matches self._generate_qconfig_map(input_root, input_graph) # match the patterns that will get quantized matches = self._find_matches(input_graph, self.modules, self.patterns) # find _inputs_ to matched nodes that are not quantized, these # have to be quantized, which requires measuring stats, # initialize an DefaultQuant object for each quants = self._find_quants(input_graph, matches) self.activation_post_process_map = dict() env = {} observed_graph = Graph() observed = set() def load_arg(a): return map_arg(a, lambda node: env[node.name]) for node in input_graph.nodes: if node.name in observed: continue def get_new_observer_name(parent_module): i = 0 def get_observer_name(i): return 'activation_post_process_' + str(i) observer_name = get_observer_name(i) while hasattr(parent_module, observer_name): i += 1 observer_name = get_observer_name(i) return observer_name root_node, _, obj, qconfig = matches.get(node.name, (None, None, None, None)) if root_node is None: env[node.name] = observed_graph.node_copy(node, load_arg) elif root_node is node: env[node.name] = observed_graph.node_copy(node, load_arg) def insert_observer(node, observer): observer_name = get_new_observer_name(input_root) setattr(input_root, observer_name, observer) self.activation_post_process_map[node.name] = observer env[node.name] = observed_graph.create_node( 'call_module', observer_name, [load_arg(node)], {}) observed.add(node.name) # don't need to insert observer for output in dynamic quantization if self.quant_type == QuantType.DYNAMIC: continue if isinstance(obj, CopyNode): assert node.op in [ 'call_module', 'call_function', 'call_method'], \ 'CopyNode of type ' + node.op + ' is not handled' def is_observed(input_arg): if isinstance(input_arg, Node): return input_arg.name in observed elif isinstance(input_arg, list): return all(map(is_observed, input_arg)) # propagate observed property from input if is_observed(node.args[0]): observed.add(node.name) elif (isinstance(obj, Add) or isinstance(obj, Mul)) and not obj.all_nodes: if node.args[0].name in observed: observed.add(node.name) elif qconfig is not None and obj.all_nodes: # observer for outputs insert_observer(node, qconfig.activation()) else: env[node.name] = observed_graph.node_copy(node, load_arg) if node.name not in observed and node.name in quants: observer_name = get_new_observer_name(input_root) _, qconfig, is_weight = quants[node.name] if qconfig is not None: self.activation_post_process_map[ node.name] = qconfig.weight( ) if is_weight else qconfig.activation() setattr(input_root, observer_name, self.activation_post_process_map[node.name]) env[node.name] = observed_graph.create_node( 'call_module', observer_name, [load_arg(node)], {}) observed.add(node.name) observed_graph.output(load_arg(input_graph.result)) return GraphModule(input_root, observed_graph)
def _prepare(self, model: GraphModule, qconfig_dict: Any, prepare_custom_config_dict: Optional[Dict[str, Any]], is_standalone_module: bool) -> GraphModule: """ standalone_module means it a submodule that is not inlined in parent module, and will be quantized separately as one unit. When we are preparing a standalone module: both input and output are observed in prepared standalone module Returns: model(GraphModule): prepared standalone module """ if prepare_custom_config_dict is None: prepare_custom_config_dict = {} self.prepare_custom_config_dict = prepare_custom_config_dict additional_quant_patterns = \ prepare_custom_config_dict.get("additional_quant_pattern", {}) self.patterns = get_combined_dict( get_default_quant_patterns(), additional_quant_patterns) flattened_qconfig_dict = get_flattened_qconfig_dict(qconfig_dict) # TODO: support regex as well propagate_qconfig_(model, flattened_qconfig_dict) if model.training: additional_qat_module_mapping = prepare_custom_config_dict.get( "additional_qat_module_mapping", {}) self._qat_swap_modules(model, additional_qat_module_mapping) self.modules = dict(model.named_modules()) convert_dict_to_ordered_dict(qconfig_dict) # map from node name to qconfig, used in _find_matches self._generate_qconfig_map(model, model.graph, qconfig_dict) # match the patterns that will get quantized standalone_module_names = prepare_custom_config_dict.get( "standalone_module_name", None) standalone_module_classes = prepare_custom_config_dict.get( "standalone_module_class", None) custom_module_classes = get_custom_module_class_keys( prepare_custom_config_dict, "float_to_observed_custom_module_class") assert self.patterns is not None matches = self._find_matches( model.graph, self.modules, self.patterns, standalone_module_names, standalone_module_classes, custom_module_classes) # find _inputs_ to matched nodes that are not quantized, these # have to be quantized, which requires measuring stats, # initialize an DefaultQuantizeHandler object for each quants: Dict[str, Tuple[DefaultQuantizeHandler, Callable]] = \ self._find_quants(model.graph, matches) self.activation_post_process_map = dict() env: Dict[Any, Any] = {} observed_graph = Graph() observed_node_names_set: Set[str] = set() def load_arg(a): return map_arg(a, lambda node: env[node.name]) # indexes for the inputs that needs to be observed standalone_module_observed_input_idxs: List[int] = [] graph_inputs = [] for node in model.graph.nodes: if node.op == 'placeholder': graph_inputs.append(node.name) get_new_observer_name = get_new_attr_name_with_prefix( 'activation_post_process_') placeholder_node_seen_cnt = 0 output_node_seen_cnt = 0 input_quantized_idxs: List[int] = self.prepare_custom_config_dict.get( "input_quantized_idxs", []) output_quantized_idxs: List[int] = self.prepare_custom_config_dict.get( "output_quantized_idxs", []) result_node : Optional[Node] = None for node in model.graph.nodes: if node.op == 'output': # If this output is hardcoded to be quantized, insert an # observer on the previous node if it does not already # exist. cur_output_node_idx = output_node_seen_cnt output_node_seen_cnt += 1 if cur_output_node_idx in output_quantized_idxs: prev_node = node.args[0] assert isinstance(prev_node, Node), \ ('hardcoding list/dict outputs to be quantized is ' + 'not supported') if prev_node.name not in observed_node_names_set: assert self.qconfig_map is not None local_qconfig = self.qconfig_map[prev_node.name] assert local_qconfig is not None, \ 'qconfig of a node before a quantized output must exist' insert_observer( prev_node, local_qconfig.activation(), model, self.activation_post_process_map, env, observed_graph, load_arg, observed_node_names_set) observed_graph.output(load_arg(node.args[0])) result_node = node continue if node.name in observed_node_names_set: continue root_node, matched_nodes, pattern, obj, qconfig = matches.get( node.name, (None, None, None, None, None)) if root_node is None: env[node.name] = observed_graph.node_copy(node, load_arg) elif root_node is node: env[node.name] = observed_graph.node_copy(node, load_arg) # index for input of custom module that needs to be observed in # parent if qconfig is not None: assert obj is not None insert_observer_for_special_module( obj, self.modules, prepare_custom_config_dict, qconfig, node) insert_observer_for_output_of_the_node( node, obj, qconfig, self.modules, model, pattern, self.activation_post_process_map, env, observed_graph, load_arg, observed_node_names_set, matched_nodes) else: env[node.name] = observed_graph.node_copy(node, load_arg) if node.op == 'placeholder': # skip adding observers at the graph input if the input is # overriden to be quantized cur_placeholder_node_idx = placeholder_node_seen_cnt placeholder_node_seen_cnt += 1 if cur_placeholder_node_idx in input_quantized_idxs: observed_node_names_set.add(node.name) continue insert_observer_for_input_arg_of_observed_node( node, observed_node_names_set, quants, model, self.activation_post_process_map, env, observed_graph, load_arg) model = GraphModule(model, observed_graph) self.save_state(model) model = mark_observed_module(model) return model
def _convert(self, model, inplace=False, debug=False, convert_custom_config_dict=None, is_standalone_module=False): """ standalone_module means it a submodule that is not inlined in parent module, and will be quantized separately as one unit. For standalone module: the inputs will be quantized by parent module, checks `_standalone_module_observed_input_idxs` of input observed model and will treat these inputs as quantized also will not dequantize the final output. Returns a quantized standalone module which accepts quantized input(if needed) and produces quantized output (if needed). """ if convert_custom_config_dict is None: convert_custom_config_dict = {} self.restore_state(model) if not inplace: model = copy.deepcopy(model) # always run weight observers in the top level forward method # for dynamic quant ops or weight only quant ops self._run_weight_observers(model) # move to cpu since we only have quantized cpu kernels model.eval().cpu() self.modules = dict(model.named_modules()) custom_module_class_mapping = convert_custom_config_dict.get( "observed_to_quantized_custom_module_class", None) matches = self._find_matches( model.graph, self.modules, self.patterns, custom_module_class_mapping=custom_module_class_mapping) quants = self._find_quants(model.graph, matches) self.quantized_graph = Graph() env = {} quant_env = {} graph_inputs = [] for node in model.graph.nodes: if node.op == 'placeholder': graph_inputs.append(node.name) def load_non_quantized(n): if n.name not in env: assert n.name in quant_env, \ 'trying to load float node but did not find node:' + n.name + \ ' in quantized or non quantized environment, env: ' + str(env) + \ ' quant_env:' + str(quant_env) env[n.name] = Proxy(quant_env[n.name]).dequantize().node return env[n.name] def load_quantized(n): if n.name not in quant_env: assert n.name in env, \ 'trying to load quantized node but did not find node:' + n.name + \ ' in float environment:' + str(env) assert n.name in quants, 'did not find quant object for node:' + n.name quant = quants[n.name][0] quant_env[n.name] = quant.convert(self, env[n.name]) return quant_env[n.name] def load_x(n): assert n.name in env or n.name in quant_env, \ 'node ' + n.name + ' does not exist in either environment' if n.name in quant_env: return quant_env[n.name] else: return env[n.name] def load_arg(quantized): """ Input: quantized, which can be None, list, boolean or tuple - if quantized is a list or tuple, then arg should be a list and the args with corresponding indexes will be quantized - if quantized is a boolean, then all args will be quantized/not quantized - if quantized is None, then we'll load the node as long as it exists Output: fn which takes arg_or_args, and loads them from the corresponding environment depending on the value of quantized. """ assert quantized is None or isinstance( quantized, (tuple, list, bool)), type(quantized) def load_arg_impl(arg_or_args): if quantized is None: return map_arg(arg_or_args, load_x) if isinstance(quantized, bool): return map_arg( arg_or_args, load_quantized if quantized else load_non_quantized) elif isinstance(quantized, (tuple, list)): assert isinstance(arg_or_args, (tuple, list)), arg_or_args loaded_args = [] # for now, we only support quantizing positional arguments for i, a in enumerate(arg_or_args): if i in quantized: loaded_args.append(map_arg(a, load_quantized)) else: loaded_args.append(map_arg(a, load_non_quantized)) return type(arg_or_args)(loaded_args) return load_arg_impl def is_quantized(node): if isinstance(node, Node): assert node.name in env or node.name in quant_env, 'Expecting node to be in the environment' # there might be nodes appearing in both environemnts, but quant_env will take # precedence if node.name in quant_env: return True elif node.name in env: return False elif isinstance(node, list): quantized = map(is_quantized, node) if all(quantized): return True elif not any(quantized): return False else: raise Exception( "partially quantized inputs in list not handled yet") for node in model.graph.nodes: if node.op == 'output': if is_standalone_module: # result are kept quantized in the quantized standalone module graph_output = map_arg(node.args[0], load_x) else: graph_output = map_arg(node.args[0], load_non_quantized) self.quantized_graph.output(graph_output) continue root_node, matched, obj, qconfig = matches.get( node.name, (None, None, None, None)) if root_node is node: if qconfig is None: result = self.quantized_graph.node_copy( node, load_non_quantized) quantized = False else: result = obj.convert( self, node, load_arg, debug=debug, convert_custom_config_dict=convert_custom_config_dict) if node.op == 'call_module' and is_observed_standalone_module( self.modules[node.target]): quantized = self.modules[ node.target]._output_is_observed else: quantized = True # Need to get correct quantized/non-quantized state for the output of CopyNode if isinstance(obj, CopyNode): assert node.op in [ 'call_module', 'call_function', 'call_method'], \ 'CopyNode of type ' + node.op + ' is not handled' quantized = is_quantized(node.args[0]) if not activation_is_statically_quantized(qconfig): quantized = False if quantized: quant_env[node.name] = result else: env[node.name] = result continue elif root_node is not None: continue # handle activation post process calls if node.op == 'call_module': if is_activation_post_process(self.modules[node.target]): observer_module = self.modules[node.target] prev_node = node.args[0] if observer_module.dtype == torch.float16: # activations are not quantized for # fp16 dynamic quantization # copy the activaiton_post_process node here # since we may need it when we insert prepack # op for weight of linear, this will be removed # later in a separate pass env[node.name] = self.quantized_graph.node_copy( node, load_non_quantized) continue if prev_node.name in quant_env: # if previous node is already quantized, we'll just remove the activation_post_process quant_env[node.name] = quant_env[prev_node.name] continue # replace activation post process with quantization ops root_module = self.modules[''] quant_env[node.name] = quantize_node( root_module, self.quantized_graph, load_non_quantized(node.args[0]), observer_module) continue if is_standalone_module and node.op == 'placeholder' and \ graph_inputs.index(node.name) in model._standalone_module_observed_input_idxs: # the node is quantized in parent module quant_env[node.name] = self.quantized_graph.node_copy( node, load_non_quantized) else: # dequantize inputs for the node that are not quantized env[node.name] = self.quantized_graph.node_copy( node, load_non_quantized) # remove activation post process act_post_process_removed_graph = Graph() env = {} def load_arg(a): return map_arg(a, lambda node: env[node.name]) for node in self.quantized_graph.nodes: if node.op == 'output': act_post_process_removed_graph.output( map_arg(node.args[0], load_arg)) continue if node.op == 'call_module' and \ is_activation_post_process(self.modules[node.target]): # remove activation post process node env[node.name] = env[node.args[0].name] else: env[node.name] = act_post_process_removed_graph.node_copy( node, load_arg) module_dict = dict(model.named_modules()) to_be_removed = [] for name, module in model.named_modules(): if is_activation_post_process( module) and not is_submodule_of_fake_quant( name, module, module_dict): to_be_removed.append(name) for n in to_be_removed: delattr(model, n) _remove_qconfig(model) model = GraphModule(model, act_post_process_removed_graph) return model
def test_typecheck_basicblock(self): class BasicBlock(torch.nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=None): super(BasicBlock, self).__init__() if norm_layer is None: norm_layer = torch.nn.BatchNorm2d if groups != 1 or base_width != 64: raise ValueError( 'BasicBlock only supports groups=1 and base_width=64') if dilation > 1: raise NotImplementedError( "Dilation > 1 not supported in BasicBlock") # Both self.conv1 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = norm_layer(planes) self.relu = torch.nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = norm_layer(planes) self.downsample = downsample self.stride = stride def forward(self, x: TensorType((2, 2, 4, 5))): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out B = BasicBlock(2, 2) ast_rewriter = RewritingTracer() graph = ast_rewriter.trace(B) traced = GraphModule(ast_rewriter.root, graph, "gm") tc = GraphTypeChecker({}, traced) tc.type_check() for n in traced.graph.nodes: if n.target == 'output': assert isinstance(n.type, TensorType) assert torch.Size(n.type.__args__) == B.forward( torch.rand(2, 2, 4, 5)).size()
def lower_to_elementwise_interpreter( orig_mod: torch.nn.Module) -> torch.nn.Module: # ===== Stage 1: Symbolic trace the module ===== mod = symbolic_trace(orig_mod) # ===== Stage 2: Lower GraphModule representation to the C++ # interpreter's instruction format ====== instructions = [] constant_idx = 0 constants = {} fn_input_names = [] target_to_name = {operator.add: "add", operator.mul: "mul"} # For each instruction, create a triple # (instruction_name : str, inputs : List[str], output : str) # to feed into the C++ interpreter for n in mod.graph.nodes: target, args, out_name = n.target, n.args, n.name assert len(n.kwargs) == 0, "kwargs currently not supported" if n.op == 'placeholder': # Placeholders specify function argument names. Save these # for later when we generate the wrapper GraphModule fn_input_names.append(target) elif n.op == 'call_function': assert target in target_to_name, "Unsupported call target " + target arg_names = [] for arg in args: if not isinstance(arg, Node): # Pull out constants. These constants will later be # fed to the interpreter C++ object via add_constant() arg_name = f'constant_{constant_idx}' constants[arg_name] = torch.Tensor( [arg] if isinstance(arg, numbers.Number ) else arg) arg_names.append(arg_name) constant_idx += 1 else: arg_names.append(arg.name) instructions.append( (target_to_name[target], arg_names, out_name)) else: raise RuntimeError('Unsupported opcode' + n.op) interpreter = torch.classes._TorchScriptTesting._ElementwiseInterpreter( ) # Load constants for k, v in constants.items(): interpreter.add_constant(k, v) # Specify names for positional input arguments interpreter.set_input_names(fn_input_names) # Load instructions interpreter.set_instructions(instructions) # Specify name for single output interpreter.set_output_name(mod.graph.result.name) # ===== Stage 3: Create a wrapper GraphModule around the interpreter ===== class WrapperModule(torch.nn.Module): def __init__(self, interpreter): super().__init__() self.interpreter = interpreter wrapper = WrapperModule(interpreter) # Create a graph that: 1) Takes function arguments 2) Invokes the interpreter # 3) Returns the speficied return value # FIXME: The following code could be greatly simplified by symbolic_trace'ing # the wrapper with a Tracer that considers the Wrapper instance a root # module, however, I can't get `__call__` exposed on TorchBind classes # without it messing up Python `hasattr` for some reason. More digging # into CPython's implementation of hasattr is probably in order... graph = torch.fx.Graph() # Add placeholders for fn inputs placeholder_nodes = [] for name in fn_input_names: placeholder_nodes.append(graph.create_node( 'placeholder', name)) # Get the interpreter object interpreter_node = graph.create_node('get_attr', 'interpreter') # Add a node to call the interpreter instance output_node = graph.create_node(op='call_method', target='__call__', args=(interpreter_node, placeholder_nodes)) # Register output graph.output(output_node) graph.lint(wrapper) # Return final GraphModule!!! return GraphModule(wrapper, graph)
def create_a_shadows_b( name_a: str, gm_a: GraphModule, name_b: str, gm_b: GraphModule, matched_subgraph_pairs: Dict[str, Tuple[NSSubgraph, NSSubgraph]], logger_cls: Callable, should_log_inputs: bool, node_type_to_io_type_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None, ) -> GraphModule: """ Creates a new GraphModule consisting of the graph of C, with the meaningful nodes of A shadowing the corresponding nodes of B. For example, Graph A: a0 -> op0_fp32 -> a1 -> op1_fp32 -> a2 Graph B: b0 -> op0_int8 -> b1 -> op1_int8 -> b2 matched_node_pairs: {'op0': (op0_fp32, op0_int8), 'op1': (op1_fp32, op1_int8)} Graph C (A shadows B): / dequant0 -> op0_fp32 -> logger_a_0 / dequant_1 -> op1_fp32 -> logger_a_1 / / b0 -------------> op0_int8 -> logger_b_0 --------------> op1_int8 -> logger_b_1 In a nutshell, this function does the following for each node pair: * copies the necessary attributes and modules from gm_a to gm_b, keeping names unique * adds a dtype cast op (dequant, quant, etc) * adds a copy of node_a in gm_b's graph * adds loggers to the outputs of node_a and node_b """ if node_type_to_io_type_map is None: node_type_to_io_type_map = get_node_type_to_io_type_map() # graph_c is the graph created from copying the nodes of graph_b and inserting # the shadows with the nodes copied from graph_a graph_c = Graph() env_c: Dict[str, Any] = {} modules = dict(gm_b.named_modules()) def load_arg(a): return map_arg(a, lambda node: env_c[node.name]) start_node_b_to_matched_subgraph_a_and_name = {} end_node_b_to_matched_subgraph_a_and_name = {} for match_name, match in matched_subgraph_pairs.items(): subgraph_a, subgraph_b = match ref_node_type_a = get_target_type_str(subgraph_a.base_op_node, gm_a) ref_node_type_b = get_target_type_str(subgraph_b.base_op_node, gm_b) start_node_b_to_matched_subgraph_a_and_name[subgraph_b.start_node] = \ (subgraph_a, match_name, ref_node_type_a, ref_node_type_b) end_node_b_to_matched_subgraph_a_and_name[subgraph_b.end_node] = \ (subgraph_a, match_name, ref_node_type_a, ref_node_type_b) for node_b in gm_b.graph.nodes: if node_b.op == 'output': graph_c.output(map_arg(node_b.args[0], load_arg)) continue # calculate the flags to determine what to do with this node node_b_is_start_node = node_b in start_node_b_to_matched_subgraph_a_and_name node_b_is_end_node = node_b in end_node_b_to_matched_subgraph_a_and_name if (node_b_is_start_node or node_b_is_end_node): if node_b_is_start_node: subgraph_a, ref_name, ref_node_type_a, ref_node_type_b = \ start_node_b_to_matched_subgraph_a_and_name[node_b] else: assert node_b_is_end_node subgraph_a, ref_name, ref_node_type_a, ref_node_type_b = \ end_node_b_to_matched_subgraph_a_and_name[node_b] # For both start_node and end_node verify that we know how to do # the dtype cast. If we do not, skip. node_input_type_a, node_output_type_a = \ get_node_first_input_and_output_type( subgraph_a.start_node, gm_a, logger_cls, node_type_to_io_type_map) node_input_type_b, node_output_type_b = \ get_node_first_input_and_output_type( node_b, gm_b, logger_cls, node_type_to_io_type_map) node_io_types_known_a_and_b = ( node_input_type_a != NodeInputOrOutputType.UNKNOWN and node_output_type_a != NodeInputOrOutputType.UNKNOWN and node_input_type_b != NodeInputOrOutputType.UNKNOWN and node_output_type_b != NodeInputOrOutputType.UNKNOWN) if not node_io_types_known_a_and_b: print( f'skipping shadow loggers for node_b: {get_target_type_str(node_b, gm_b)}' + f', start_node_a: {get_target_type_str(subgraph_a.start_node, gm_a)}' + ', unknown dtype cast') env_c[node_b.name] = graph_c.node_copy(node_b, load_arg) continue # If we are shadowing from fp32 to int8, we need to insert # quantize_per_tensor call with qparams from the previous node. # Only do this if we are able to infer these qparams from the graph. if (node_input_type_a == NodeInputOrOutputType.INT8 and node_input_type_b == NodeInputOrOutputType.FP32): node_a_input_qparams = get_node_input_qparams( subgraph_a.start_node, gm_a, node_type_to_io_type_map) if not node_a_input_qparams: print( f'skipping shadow loggers for node_b: {get_target_type_str(node_b, gm_b)}' + f', start_node_a: {get_target_type_str(subgraph_a.start_node, gm_a)}' + ', unknown input qparams') env_c[node_b.name] = graph_c.node_copy(node_b, load_arg) continue fqn_base_a = _maybe_get_fqn(subgraph_a.base_op_node, gm_a) fqn_base_b = _maybe_get_fqn(subgraph_b.base_op_node, gm_b) if node_b_is_start_node: # if necessary, log the input of node_c if should_log_inputs: if isinstance(node_b.args[0], Node): prev_node_c = env_c[node_b.args[0].name] env_c[prev_node_c.name] = _insert_logger_after_node( prev_node_c, gm_b, logger_cls, '_ns_logger_b_inp_', node_b.name, name_b, ref_name, ref_node_type_b, NSSingleResultValuesType.NODE_INPUT.value, index_within_arg=0, index_of_arg=0, fqn=fqn_base_b) elif isinstance(node_b.args[0], list): # first, save the prev_node instances, because they # will be overwritten in the env after the first logger # is added prev_node_c_list = [ env_c[arg.name] for arg in node_b.args[0] ] for arg_idx, arg in enumerate(node_b.args[0]): prev_node_c = prev_node_c_list[arg_idx] env_c[ prev_node_c.name] = _insert_logger_after_node( prev_node_c, gm_b, logger_cls, '_ns_logger_b_inp_', node_b.name, name_b, ref_name, ref_node_type_b, NSSingleResultValuesType.NODE_INPUT.value, index_within_arg=arg_idx, index_of_arg=0, fqn=fqn_base_b) else: # logging of inputs which are not lists is not supported yet raise AssertionError( f"type {type(node_b.args[0])} is not handled yet") # subgraph so far: # # (prev_node_c)+ -> (logger_c_input)? # Note: this if statement is always True, spelling it out to clarify code # intent. if node_b_is_start_node or node_b_is_end_node: # ensure env_c is populated with base node env_c[node_b.name] = graph_c.node_copy(node_b, load_arg) node_c = env_c[node_b.name] # after this point, # # node_a is the original node from graph_a, with parent module gm_a # node_b is the original node from graph_b, with parent module gm_b # node_c is the copy of node_b in graph_c # # subgraph so far: # # (prev_node_c)+ -> (logger_c_input)? -> node_start_c if node_b_is_start_node: # cast dtype from the dtype of node_c's input to the dtype of # node_a's input (dequant, etc) prev_node_c = node_c.args[0] if should_log_inputs: # skip the input logger when inserting a dtype cast if isinstance(prev_node_c, Node): prev_node_c = prev_node_c.args[0] elif isinstance(prev_node_c, list): prev_node_c = [arg.args[0] for arg in prev_node_c] dtype_cast_node = _insert_dtype_cast_after_node( subgraph_a.start_node, node_c, prev_node_c, gm_a, gm_b, graph_c, node_b.name + '_dtype_cast_', logger_cls, node_type_to_io_type_map) # note: not inserting to env_c because all nodes which use the dtype # casts are copied from graph_a # # subgraph so far: # # (dtype_cast_node)+ # / # (prev_node_c)+ -> (logger_c_input)? -> node_start_c # if input logging is enabled, log the input to the subgraph if should_log_inputs: # TODO: explain this ref_node_name = '' if isinstance(dtype_cast_node, Node): dtype_cast_node = _insert_logger_after_node( dtype_cast_node, gm_b, logger_cls, '_ns_logger_a_inp_', ref_node_name, name_a, ref_name, ref_node_type_a, NSSingleResultValuesType.NODE_INPUT.value, index_within_arg=0, index_of_arg=0, fqn=fqn_base_a) input_logger: Union[Node, List[Node]] = dtype_cast_node else: assert isinstance(dtype_cast_node, list) new_loggers = [] for dtype_cast_idx, dtype_cast_node_inner in enumerate( dtype_cast_node): dtype_cast_logger = _insert_logger_after_node( dtype_cast_node_inner, gm_b, logger_cls, '_ns_logger_a_inp_', ref_node_name, name_a, ref_name, ref_node_type_a, NSSingleResultValuesType.NODE_INPUT.value, index_within_arg=dtype_cast_idx, index_of_arg=0, fqn=fqn_base_a) new_loggers.append(dtype_cast_logger) dtype_cast_node = new_loggers input_logger = dtype_cast_node # subgraph so far: # # (dtype_cast_node)+ -> (logger_a_input)? # / # prev_node_c -> (logger_c_input)? -> node_start_c # hook up the new mod_a copy to be in the graph, receiving the # same inputs as mod_b does, with dtype cast to match a # Some ops, such as LSTMs, have two non-param inputs. If we have # such an op, pass the second param as well. Note: dtype casting # for the second param is not implemented yet, it can be added # later if there is a use case. node_c_second_non_param_arg = None num_non_param_args_node_a = get_number_of_non_param_args( subgraph_a.start_node, gm_a) if num_non_param_args_node_a == 2: node_c_second_non_param_arg = node_c.args[1] node_a_shadows_c = _insert_copy_of_subgraph_a_after_input_node_c( dtype_cast_node, node_c_second_non_param_arg, subgraph_a, gm_a, gm_b, node_c.name + '_shadow_copy_') env_c[node_a_shadows_c.name] = node_a_shadows_c # subgraph so far: # # dtype_cast_node -> (logger_a_input)? -> subgraph_a_copy(args/kwargs not shown) # / # (prev_node_c)+ -> (logger_c_input)? -> node_start_c if should_log_inputs: # When we created the input logger, we left the ref_node_name # as an empty string, because the subgraph copy did not exist # yet. Now that the subgraph copy exists, we modify this name # to its true value. # Note: the alternative to this is to create the input logger # after creating the subgraph, which is slightly more # complicated. This is the lesser of two evils. # input_logger = env_c[dtype_cast_node.name] # Find the first node in the subgraph cur_node = node_a_shadows_c while cur_node.args[0] != input_logger: cur_node = cur_node.args[0] # type: ignore[assignment] if isinstance(input_logger, Node): input_logger_mod = getattr(gm_b, input_logger.name) input_logger_mod.ref_node_name = cur_node.name else: assert isinstance(input_logger, list) for input_logger_inner in input_logger: input_logger_mod = getattr(gm_b, input_logger_inner.name) input_logger_mod.ref_node_name = cur_node.name # hook up a logger to the mod_a copy env_c[node_a_shadows_c.name] = _insert_logger_after_node( env_c[node_a_shadows_c.name], gm_b, logger_cls, '_ns_logger_a_', node_a_shadows_c.name, name_a, ref_name, ref_node_type_a, NSSingleResultValuesType.NODE_OUTPUT.value, index_within_arg=0, index_of_arg=0, fqn=fqn_base_a) # subgraph so far: # # dtype_cast_node -> (logger_a_input)? -> subgraph_a_copy -> logger_a # / # (prev_node_c)+ -> (logger_c_input)? -> node_start_c if node_b_is_end_node: # hook up a logger to the mod_b copy env_c[node_b.name] = _insert_logger_after_node( env_c[node_b.name], gm_b, logger_cls, '_ns_logger_b_', node_b.name, name_b, ref_name, ref_node_type_b, NSSingleResultValuesType.NODE_OUTPUT.value, index_within_arg=0, index_of_arg=0, fqn=fqn_base_b) # subgraph so far: # # dtype_cast_node -> (logger_a_input)? -> subgraph_a_copy -> logger_a # / # (prev_node_c+) -> (logger_c_input)? -> node_start_c -> ... -> node_end_c -> logger_c # # Note: node_start_c may be the same node as node_end_c, or they # may have nodes inbetween. else: env_c[node_b.name] = graph_c.node_copy(node_b, load_arg) gm_c = GraphModule(gm_b, graph_c) return gm_c
def _convert(self, model: GraphModule, debug: bool = False, convert_custom_config_dict: Dict[str, Any] = None, is_standalone_module: bool = False) -> GraphModule: """ standalone_module means it a submodule that is not inlined in parent module, and will be quantized separately as one unit. Returns a quantized standalone module which accepts float input and produces float output. """ if convert_custom_config_dict is None: convert_custom_config_dict = {} self.restore_state(model) # always run weight observers in the top level forward method # for dynamic quant ops or weight only quant ops self._run_weight_observers(model) # move to cpu since we only have quantized cpu kernels model.eval().cpu() self.modules = dict(model.named_modules()) custom_module_classes = get_custom_module_class_keys( convert_custom_config_dict, "observed_to_quantized_custom_module_class") assert self.patterns is not None matches = self._find_matches( model.graph, self.modules, self.patterns, custom_module_classes=custom_module_classes) quants: Dict[str, Tuple[DefaultQuantizeHandler, Callable]] = \ self._find_quants(model.graph, matches) self.quantized_graph = Graph() env: Dict[str, Node] = {} quant_env: Dict[str, Node] = {} graph_inputs: List[str] = [] for node in model.graph.nodes: if node.op == 'placeholder': graph_inputs.append(node.name) def load_non_quantized(n: Node) -> Node: if n.name not in env: assert n.name in quant_env, \ 'trying to load float node but did not find ' + \ 'node:' + n.name + \ ' in quantized or non quantized environment, env: ' + \ str(env) + ' quant_env:' + str(quant_env) env[n.name] = Proxy(quant_env[n.name]).dequantize().node return env[n.name] def load_quantized(n: Node) -> Node: assert n.name in quant_env, \ 'trying to load quantized node but did not find node:' + \ n.name + ' in quant environment:' + str(quant_env) return quant_env[n.name] def load_x(n: Node) -> Node: assert n.name in env or n.name in quant_env, \ 'node ' + n.name + ' does not exist in either environment' if n.name in quant_env: return quant_env[n.name] else: return env[n.name] def load_arg(quantized: Optional[Union[List[Any], bool, Tuple[Any, ...]]] ) -> Callable[[Node], Argument]: """ Input: quantized, which can be None, list, boolean or tuple - if quantized is a list or tuple, then arg should be a list and the args with corresponding indexes will be quantized - if quantized is a boolean, then all args will be quantized/not quantized - if quantized is None, then we'll load the node as long as it exists Output: fn which takes arg_or_args, and loads them from the corresponding environment depending on the value of quantized. """ assert quantized is None or \ isinstance(quantized, (tuple, list, bool)), type(quantized) def load_arg_impl(arg_or_args): if quantized is None: return map_arg(arg_or_args, load_x) if isinstance(quantized, bool): return map_arg( arg_or_args, load_quantized if quantized else load_non_quantized) elif isinstance(quantized, (tuple, list)): assert isinstance(arg_or_args, (tuple, list)), arg_or_args loaded_args = [] # for now, we only support quantizing positional arguments for i, a in enumerate(arg_or_args): if i in quantized: loaded_args.append(map_arg(a, load_quantized)) else: loaded_args.append(map_arg(a, load_non_quantized)) return type(arg_or_args)(loaded_args) return load_arg_impl def node_arg_is_quantized(node_arg: Any) -> bool: if isinstance(node_arg, Node): assert node_arg.name in env or node_arg.name in quant_env, \ 'Expecting node_arg to be in the environment' # there might be nodes appearing in both environemnts, but # quant_env will take precedence if node_arg.name in quant_env: return True elif node_arg.name in env: return False else: return False elif isinstance(node_arg, list): quantized = map(node_arg_is_quantized, node_arg) if all(quantized): return True elif not any(quantized): return False else: raise Exception( "partially quantized inputs in list not handled yet") else: return False def is_output_quantized(node: Node, obj: QuantizeHandler) -> bool: """ Check if output node is quantized or not """ assert self.modules is not None # by default the output is expected to be quantized quantized = True # Need to get correct quantized/non-quantized state for the output # of CopyNode if type(obj) in [ CopyNode, FixedQParamsOpQuantizeHandler ]: assert node.op in [ 'call_module', 'call_function', 'call_method'], \ 'CopyNode of type ' + node.op + ' is not handled' quantized = node_arg_is_quantized(node.args[0]) if not activation_is_statically_quantized(qconfig) or \ not input_output_observed(obj): quantized = False return quantized def insert_quantize_node(node: Node) -> None: """ Given a activation_post_process module call node, insert a quantize node""" assert self.modules is not None assert isinstance(node.target, str) observer_module = self.modules[node.target] prev_node = node.args[0] if observer_module.dtype == torch.float16: # activations are not quantized for # fp16 dynamic quantization # copy the activaiton_post_process node here # since we may need it when we insert prepack # op for weight of linear, this will be removed # later in a separate pass env[node.name] = self.quantized_graph.node_copy( node, load_non_quantized) elif isinstance(prev_node, Node) and prev_node.name in quant_env: # if previous node is already quantized, we'll just remove the # activation_post_process quant_env[node.name] = quant_env[prev_node.name] else: # replace activation post process with quantization ops root_module = self.modules[""] assert isinstance(node.args[0], Node) quant_env[node.name] = quantize_node( root_module, self.quantized_graph, load_non_quantized(node.args[0]), observer_module) # additional state to override inputs to be quantized, if specified # by the user placeholder_node_seen_cnt = 0 output_node_seen_cnt = 0 input_quantized_idxs: List[int] = self.prepare_custom_config_dict.get( "input_quantized_idxs", []) output_quantized_idxs: List[int] = self.prepare_custom_config_dict.get( "output_quantized_idxs", []) for node in model.graph.nodes: if node.op == 'output': cur_output_node_idx = output_node_seen_cnt output_node_seen_cnt += 1 if cur_output_node_idx in output_quantized_idxs: # Result are kept quantized if the user specified the # output_quantized_idxs override. graph_output = map_arg(node.args[0], load_x) else: graph_output = map_arg(node.args[0], load_non_quantized) self.quantized_graph.output(graph_output) continue root_node, matched, matched_pattern, obj, qconfig = \ matches.get(node.name, (None, None, None, None, None)) if root_node is node: if qconfig is None: result = self.quantized_graph.node_copy( node, load_non_quantized) quantized = False else: assert obj is not None is_standalone_module_node = ( node.op == 'call_module' and is_observed_standalone_module( self.modules[node.target]) # type: ignore ) result = obj.convert( self, node, load_arg, debug=debug, convert_custom_config_dict=convert_custom_config_dict) if is_standalone_module_node: quantized = False else: quantized = is_output_quantized(node, obj) if quantized: quant_env[node.name] = result else: env[node.name] = result continue elif root_node is not None: continue # handle activation post process calls if node.op == 'call_module' and \ is_activation_post_process(self.modules[node.target]): insert_quantize_node(node) elif node.op == 'placeholder': cur_placeholder_node_idx = placeholder_node_seen_cnt placeholder_node_seen_cnt += 1 if cur_placeholder_node_idx in input_quantized_idxs: quant_env[node.name] = \ self.quantized_graph.node_copy(node, load_non_quantized) else: env[node.name] = \ self.quantized_graph.node_copy(node, load_non_quantized) else: # copy quantized or non-quantized node env[node.name] = \ self.quantized_graph.node_copy(node, load_non_quantized) # remove activation post process act_post_process_removed_graph = Graph() env = {} def load_arg_simple(a: Argument) -> Argument: return map_arg(a, lambda node: env[node.name]) for node in self.quantized_graph.nodes: if node.op == 'output': act_post_process_removed_graph.output( map_arg(node.args[0], load_arg_simple)) continue if node.op == 'call_module' and \ is_activation_post_process(self.modules[node.target]): # remove activation post process node env[node.name] = env[node.args[0].name] else: env[node.name] = act_post_process_removed_graph.node_copy( node, load_arg_simple) # removes qconfig and activation_post_process modules _remove_qconfig(model) model = GraphModule(model, act_post_process_removed_graph) return model
def add_loggers_to_model( gm: GraphModule, node_to_instrument_inputs_to_ref_node_name: Dict[Node, Tuple[str, str]], node_to_instrument_outputs_to_ref_node_name: Dict[Node, Tuple[str, str]], logger_cls: Callable, model_name: str, ) -> GraphModule: """ Takes the graph of gm, adds loggers to the output of each node in nodes_to_instrument. Returns a GraphModule with the new graph. """ new_graph = Graph() env: Dict[str, Any] = {} modules = dict(gm.named_modules()) def load_arg(a): return map_arg(a, lambda node: env[node.name]) for node in gm.graph.nodes: if node.op == 'output': new_graph.output(map_arg(node.args[0], load_arg)) continue if ((node in node_to_instrument_inputs_to_ref_node_name) or (node in node_to_instrument_outputs_to_ref_node_name)): fqn = _maybe_get_fqn(node, gm) if node in node_to_instrument_inputs_to_ref_node_name: ref_name, ref_node_type = node_to_instrument_inputs_to_ref_node_name[ node] # Ops such add and mul are special because either # one or two of the first two arguments can be tensors, # and if one argument is a tensor it can be first or # second (x + 1 versus 1 + x). arg_indices_to_log = get_arg_indices_of_inputs_to_log(node) for node_arg_idx in arg_indices_to_log: node_arg = node.args[node_arg_idx] if type(node_arg) == Node: # create a single input logger prev_node = env[node_arg.name] env[node_arg.name] = _insert_logger_after_node( prev_node, gm, logger_cls, '_ns_logger_', node.name, model_name, ref_name, ref_node_type, NSSingleResultValuesType.NODE_INPUT.value, index_within_arg=0, index_of_arg=node_arg_idx, fqn=fqn) elif type( node_arg ) == torch.fx.immutable_collections.immutable_list: # create N input loggers, one for each node for arg_idx, arg in enumerate(node_arg): prev_node = env[arg.name] env[prev_node.name] = _insert_logger_after_node( prev_node, gm, logger_cls, '_ns_logger_', node.name, model_name, ref_name, ref_node_type, NSSingleResultValuesType.NODE_INPUT.value, index_within_arg=arg_idx, index_of_arg=node_arg_idx, fqn=fqn) else: pass # ensure env is populated with base node # Note: runs for both inputs and outputs env[node.name] = new_graph.node_copy(node, load_arg) if node in node_to_instrument_outputs_to_ref_node_name: ref_name, ref_node_type = node_to_instrument_outputs_to_ref_node_name[ node] # add the logger after the base node env[node.name] = _insert_logger_after_node( env[node.name], gm, logger_cls, '_ns_logger_', node.name, model_name, ref_name, ref_node_type, NSSingleResultValuesType.NODE_OUTPUT.value, index_within_arg=0, index_of_arg=0, fqn=fqn) else: env[node.name] = new_graph.node_copy(node, load_arg) new_gm = GraphModule(gm, new_graph) return new_gm
def convert(self, observed, inplace=False, debug=False): assert self.activation_post_process_map is not None # move to cpu since we only have quantized cpu kernels observed.eval().cpu() observed_root = observed.root observed_graph = observed.graph if not inplace: observed_root = copy.deepcopy(observed_root) self.modules = dict(observed_root.named_modules()) matches = self._find_matches(observed.graph, self.modules, self.patterns) quants = self._find_quants(observed.graph, matches) self.quantized_graph = Graph() env = {} quant_env = {} def load_non_quantized(n): if n.name not in env: assert n.name in quant_env, \ 'trying to load float node but did not find node:' + n.name + \ ' in quantized environment:' + str(quant_env) env[n.name] = Proxy(quant_env[n.name]).dequantize().node return env[n.name] def load_quantized(n): if n.name not in quant_env: assert n.name in env, \ 'trying to load quantized node but did not find node:' + n.name + \ ' in float environment:' + str(env) assert n.name in quants, 'did not find quant object for node:' + n.name quant = quants[n.name][0] quant_env[n.name] = quant.convert(self, env[n.name]) return quant_env[n.name] def load_x(n): assert n.name in env or n.name in quant_env, \ 'node ' + n.name + ' does not exist in either of the environment' if n.name in quant_env: return quant_env[n.name] else: return env[n.name] def load_arg(quantized): """ if quantized is a list, then arg should be a list and the args with corresponding indexes will be quantized if quantized is a boolean, then all args will be quantized/not quantized if quantized is None, then we'll load the node as long as it exists """ assert quantized is None or isinstance( quantized, (tuple, list, bool)), type(quantized) def load_arg_impl(arg): if quantized is None: return map_arg(arg, load_x) if isinstance(quantized, bool): return map_arg( arg, load_quantized if quantized else load_non_quantized) elif isinstance(quantized, (tuple, list)): assert isinstance(arg, (tuple, list)), arg loaded_arg = [] # for now, we only support quantizing positional arguments for i, a in enumerate(arg): if i in quantized: loaded_arg.append(map_arg(a, load_quantized)) else: loaded_arg.append(map_arg(a, load_non_quantized)) return type(arg)(loaded_arg) return load_arg_impl def is_quantized(node): if isinstance(node, Node): assert node.name in env or node.name in quant_env, 'Expecting node to be in the environment' # there might be nodes appearing in both environemnts, but quant_env will take # precedence if node.name in quant_env: return True elif node.name in env: return False elif isinstance(node, list): quantized = map(is_quantized, node) if all(quantized): return True elif not any(quantized): return False else: raise Exception( "partially quantized inputs in list not handled yet") for node in observed_graph.nodes: root_node, matched, obj, qconfig = matches.get( node.name, (None, None, None, None)) if root_node is node: result = obj.convert(self, node, load_arg) quantized = True # Need to get correct quantized/non-quantized state for the output of CopyNode if isinstance(obj, CopyNode): assert node.op in [ 'call_module', 'call_function', 'call_method'], \ 'CopyNode of type ' + node.op + ' is not handled' quantized = is_quantized(node.args[0]) if self.quant_type == QuantType.DYNAMIC: quantized = False if quantized: quant_env[node.name] = result else: env[node.name] = result continue elif root_node is not None: continue # handle activation post process calls if node.op == 'call_module': if node.target.split('.')[-1].startswith( 'activation_post_process_'): observer_module = self.modules[node.target] prev_node = node.args[0] if prev_node.name in quant_env: # if previous node is already quantized, we'll just remove the activation_post_process quant_env[node.name] = quant_env[prev_node.name] continue # replace activation post process with quantization ops parent_name = '' scale, zero_point = observer_module.calculate_qparams() # TODO: per channel scale = float(scale) zero_point = int(zero_point) dtype = observer_module.dtype qparams = { '_scale_': scale, '_zero_point_': zero_point, '_dtype_': dtype } i = 0 def noattr(module, qparams, i): for name in qparams.keys(): if hasattr(module, name + str(i)): return False return True def get_next_i(module, qparams): i = 0 while not noattr(module, qparams, i): i += 1 return i parent_module = self.modules[parent_name] i = get_next_i(parent_module, qparams) inputs = [load_non_quantized(node.args[0])] for key, value in qparams.items(): setattr(parent_module, key + str(i), value) qparam_full_path = key + str(i) if parent_name: qparam_full_path = parent_name + '.' + qparam_full_path inputs.append( self.quantized_graph.create_node( 'get_param', qparam_full_path)) quant_env[node.name] = self.quantized_graph.create_node( 'call_function', torch.quantize_per_tensor, inputs, {}) continue # dequantize inputs for the node that are not quantized env[node.name] = self.quantized_graph.node_copy( node, load_non_quantized) self.quantized_graph.output(load_non_quantized(observed_graph.result)) to_be_removed = [] for name, _ in observed_root.named_modules(): if name.split('.')[-1].startswith('activation_post_process_'): to_be_removed.append(name) for n in to_be_removed: delattr(observed_root, n) return GraphModule(observed_root, self.quantized_graph)
def _prepare_fx(model: torch.nn.Module, qconfig_dict: Any, prepare_custom_config_dict: Optional[Dict[str, Any]] = None, equalization_qconfig_dict: Optional[Dict[str, Any]] = None, backend_config_dict: Optional[Dict[str, Any]] = None, is_standalone_module: bool = False) -> ObservedGraphModule: r""" Internal helper function for prepare_fx Args: `model`, `qconfig_dict`, `prepare_custom_config_dict`, `equalization_qonfig_dict`: see docs for :func:`~torch.quantization.prepare_fx` `is_standalone_module`: a boolean flag indicates whether we are quantizing a standalone module or not, a standalone module is a submodule of the parent module that is not inlined in the forward graph of the parent module, the way we quantize standalone module is described in: :func:`~torch.quantization._prepare_standalone_module_fx` """ if prepare_custom_config_dict is None: prepare_custom_config_dict = {} if equalization_qconfig_dict is None: equalization_qconfig_dict = {} check_is_valid_qconfig_dict(qconfig_dict) check_is_valid_prepare_custom_config_dict(prepare_custom_config_dict) check_is_valid_qconfig_dict(equalization_qconfig_dict) skipped_module_names = prepare_custom_config_dict.get( "non_traceable_module_name", []) skipped_module_classes = prepare_custom_config_dict.get( "non_traceable_module_class", []) # swap FloatFunctional with FXFloatFunctional _swap_ff_with_fxff(model) # symbolically trace the model if not is_standalone_module: # standalone module and custom module config are applied in top level module standalone_module_name_configs = prepare_custom_config_dict.get( "standalone_module_name", []) skipped_module_names += [ config[0] for config in standalone_module_name_configs ] standalone_module_class_configs = prepare_custom_config_dict.get( "standalone_module_class", []) skipped_module_classes += [ config[0] for config in standalone_module_class_configs ] float_custom_module_classes = get_custom_module_class_keys( prepare_custom_config_dict, "float_to_observed_custom_module_class") skipped_module_classes += float_custom_module_classes preserved_attributes = prepare_custom_config_dict.get( "preserved_attributes", []) tracer = QuantizationTracer(skipped_module_names, skipped_module_classes) graph_module = GraphModule(model, tracer.trace(model)) for attr_name in preserved_attributes: setattr(graph_module, attr_name, getattr(model, attr_name)) graph_module = _fuse_fx(graph_module, prepare_custom_config_dict) prepared = prepare(graph_module, qconfig_dict, tracer.node_name_to_scope, prepare_custom_config_dict=prepare_custom_config_dict, equalization_qconfig_dict=equalization_qconfig_dict, is_standalone_module=is_standalone_module) for attr_name in preserved_attributes: setattr(prepared, attr_name, getattr(model, attr_name)) return prepared