def add_adjust_padding_nodes(bitwidth_graph: nx.DiGraph,
model: NNCFNetwork) -> nx.DiGraph():
# pylint:disable=protected-access
NewNodeArgs = namedtuple('NewNodeArgs',
('node_key', 'attr', 'parent_node_key'))
nncf_graph = model.get_graph()
args = []
for node_key in bitwidth_graph.nodes:
node = nncf_graph.get_node_by_key(node_key)
module = model.get_containing_module(node.node_name)
if isinstance(module, NNCFConv2d):
adjust_padding_ops = filter(
lambda x: isinstance(x, UpdatePaddingValue),
module.pre_ops.values())
for _ in adjust_padding_ops:
new_node_key = f'{node_key}_apad'
attr = dict(type='',
label='adjust_padding_value',
style='filled',
color='yellow')
args.append(NewNodeArgs(new_node_key, attr, node_key))
for arg in args:
bitwidth_graph.add_node(arg.node_key, **arg.attr)
bitwidth_graph.add_edge(arg.node_key, arg.parent_node_key)
return bitwidth_graph
def check_model_graph(compressed_model: NNCFNetwork, ref_dot_file_name: str, ref_dot_file_directory: str):
if torch.cuda.is_available():
compressed_model.to('cuda')
compressed_model.do_dummy_forward()
# internal wrapped model is still in eval mode, switch to the train mode to make sure training graph is ok
compressed_model.train()
compressed_model.rebuild_graph()
compressed_model.do_dummy_forward()
check_graph(compressed_model.get_graph(), ref_dot_file_name, ref_dot_file_directory)
def __init__(self, algo_ctrl: QuantizationController, model: NNCFNetwork,
groups_of_adjacent_quantizers: GroupsOfAdjacentQuantizers,
add_flops=False):
# pylint:disable=too-many-branches
# pylint:disable=too-many-statements
nncf_graph = model.get_graph()
self._nx_graph = nncf_graph.get_graph_for_structure_analysis()
if add_flops:
flops_per_module = model.get_flops_per_module()
flops_vs_node_group = defaultdict(set) # type: Dict[int, Tuple[int, Set[NNCFNode]]]
for idx, module_node_name_and_flops in enumerate(flops_per_module.items()):
module_node_name, flops = module_node_name_and_flops
node_set = set(nncf_graph.get_op_nodes_in_scope(nncf_graph.get_scope_by_node_name(module_node_name)))
flops_vs_node_group[idx] = (flops, node_set)
grouped_mode = bool(groups_of_adjacent_quantizers)
for node_key in nncf_graph.get_all_node_keys():
node = nncf_graph.get_node_by_key(node_key)
color = ''
operator_name = node.node_type
module = model.get_containing_module(node.node_name)
if isinstance(module, NNCFConv2d):
color = 'lightblue'
if module.groups == module.in_channels and module.in_channels > 1:
operator_name = 'DW_Conv2d'
color = 'purple'
kernel_size = 'x'.join(map(str, module.kernel_size))
operator_name += f'_k{kernel_size}'
padding_values = set(module.padding)
padding_enabled = len(padding_values) >= 1 and padding_values.pop()
if padding_enabled:
operator_name += '_PAD'
if add_flops:
matches = [f_nodes_tpl for idx, f_nodes_tpl in flops_vs_node_group.items()
if node in f_nodes_tpl[1]]
assert len(matches) == 1
flops, affected_nodes = next(iter(matches))
operator_name += f'_FLOPS:{str(flops)}'
if len(affected_nodes) > 1:
node_ids = sorted([n.node_id for n in affected_nodes])
operator_name += "(shared among nodes {})".format(",".join(
[str(node_id) for node_id in node_ids]))
operator_name += '_#{}'.format(node.node_id)
target_node_to_draw = self._nx_graph.nodes[node_key]
target_node_to_draw['label'] = operator_name
target_node_to_draw['style'] = 'filled'
if color:
target_node_to_draw['color'] = color
non_weight_quantizers = algo_ctrl.non_weight_quantizers
bitwidth_color_map = {2: 'purple', 4: 'red', 8: 'green', 6: 'orange'}
for quantizer_id, quantizer_info in non_weight_quantizers.items():
self._paint_activation_quantizer_node(nncf_graph, quantizer_id,
quantizer_info, bitwidth_color_map,
groups_of_adjacent_quantizers)
for wq_id, wq_info in algo_ctrl.weight_quantizers.items():
nodes = [nncf_graph.get_node_by_name(tp.target_node_name)
for tp in wq_info.affected_insertions]
if not nodes:
raise AttributeError('Failed to get affected nodes for quantized module node: {}'.format(
wq_id.target_node_name))
preds = [nncf_graph.get_previous_nodes(node) for node in nodes]
wq_nodes = []
for pred_list in preds:
for pred_node in pred_list:
if 'UpdateWeight' in pred_node.node_name:
wq_nodes.append(pred_node)
assert len(wq_nodes) == 1
node = wq_nodes[0]
node_id = node.node_id
key = nncf_graph.get_node_key_by_id(node_id)
nx_node_to_draw_upon = self._nx_graph.nodes[key]
quantizer = wq_info.quantizer_module_ref
bitwidths = quantizer.num_bits
nx_node_to_draw_upon['label'] = 'WFQ_[{}]_#{}'.format(quantizer.get_quantizer_config(), str(node_id))
if grouped_mode:
group_id_str = 'UNDEFINED'
group_id = groups_of_adjacent_quantizers.get_group_id_for_quantizer(wq_id)
if group_id is None:
nncf_logger.error('No group for weight quantizer for: {}'.format(wq_id))
else:
group_id_str = str(group_id)
nx_node_to_draw_upon['label'] += '_G' + group_id_str
nx_node_to_draw_upon['color'] = bitwidth_color_map[bitwidths]
nx_node_to_draw_upon['style'] = 'filled'
def create_compressed_model(model: Module,
config: NNCFConfig,
compression_state: Optional[Dict[str, Any]] = None,
dummy_forward_fn: Callable[[Module], Any] = None,
wrap_inputs_fn: Callable[[Tuple, Dict], Tuple[Tuple, Dict]] = None,
wrap_outputs_fn: Callable[[Tuple, Dict], Tuple[Tuple, Dict]] = None,
dump_graphs=True) \
-> Tuple[CompressionAlgorithmController, NNCFNetwork]:
"""
The main function used to produce a model ready for compression fine-tuning from an original PyTorch
model and a configuration object.
dummy_forward_fn
:param model: The original model. Should have its parameters already loaded from a checkpoint or another
source.
:param config: A configuration object used to determine the exact compression modifications to be applied
to the model
:param compression_state: representation of the entire compression state to unambiguously restore
the compressed model. Includes builder and controller states.
:param dummy_forward_fn: if supplied, will be used instead of a *forward* function call to build
the internal graph representation via tracing. Specifying this is useful when the original training pipeline
has special formats of data loader output or has additional *forward* arguments other than input tensors.
Otherwise, the *forward* call of the model during graph tracing will be made with mock tensors according
to the shape specified in the config object. The dummy_forward_fn code MUST contain calls to nncf.nncf_model_input
functions made with each compressed model input tensor in the underlying model's args/kwargs tuple, and these
calls should be exactly the same as in the wrap_inputs_fn function code (see below); if dummy_forward_fn is
specified, then wrap_inputs_fn also must be specified.
:param wrap_inputs_fn: if supplied, will be used on the module's input arguments during a regular, non-dummy
forward call before passing the inputs to the underlying compressed model. This is required if the model's input
tensors that are important for compression are not supplied as arguments to the model's forward call directly, but
instead are located in a container (such as list), and the model receives the container as an argument.
wrap_inputs_fn should take as input two arguments - the tuple of positional arguments to the underlying
model's forward call, and a dict of keyword arguments to the same. The function should wrap each tensor among the
supplied model's args and kwargs that is important for compression (e.g. quantization) with an nncf.nncf_model_input
function, which is a no-operation function and marks the tensors as inputs to be traced by NNCF in the internal
graph representation. Output is the tuple of (args, kwargs), where args and kwargs are the same as were supplied in
input, but each tensor in the original input. Must be specified if dummy_forward_fn is specified.
:param dump_graphs: Whether or not should also dump the internal graph representation of the
original and compressed models in the .dot format into the log directory.
:return: A controller for the compression algorithm (or algorithms, in which case the controller
is an instance of CompositeCompressionController) and the model ready for compression parameter training wrapped
as an object of NNCFNetwork."""
if dummy_forward_fn is not None and wrap_inputs_fn is None:
raise ValueError(
"A custom dummy forward function was specified, but the corresponding input wrapping function "
"was not. In case a custom dummy forward function is specified for purposes of NNCF graph "
"building, then the wrap_inputs_fn parameter MUST also be specified and be consistent with "
"the input wrapping done in dummy_forward_fn.")
is_legacy_model_state_dict = compression_state is not None and \
BaseController.BUILDER_STATE not in compression_state and \
BaseController.CONTROLLER_STATE not in compression_state
maybe_convert_legacy_names_in_compress_state(compression_state)
# Compress model that will be deployed for the inference on target device. No need to compress parts of the
# model that are used on training stage only (e.g. AuxLogits of Inception-v3 model) or unused modules with weights.
# As a consequence, no need to care about spoiling BN statistics, as there're disabled in eval mode.
model.eval()
if dump_graphs:
if dummy_forward_fn is None:
input_info_list = create_input_infos(config)
graph_builder = GraphBuilder(
custom_forward_fn=create_dummy_forward_fn(
input_info_list, with_input_tracing=True))
else:
graph_builder = GraphBuilder(custom_forward_fn=dummy_forward_fn)
if is_main_process():
graph = graph_builder.build_graph(model)
graph.visualize_graph(
osp.join(config.get("log_dir", "."), "original_graph.dot"))
set_debug_log_dir(config.get("log_dir", "."))
input_info_list = create_input_infos(config)
scopes_without_shape_matching = config.get('scopes_without_shape_matching',
[])
ignored_scopes = config.get('ignored_scopes')
target_scopes = config.get('target_scopes')
original_model_accuracy = None
if is_accuracy_aware_training(config):
if config.has_extra_struct(ModelEvaluationArgs):
evaluation_args = config.get_extra_struct(ModelEvaluationArgs)
with torch.no_grad():
original_model_accuracy = evaluation_args.eval_fn(model)
nncf_logger.info("Non-compressed model accuracy = {}".format(
original_model_accuracy))
compressed_model = NNCFNetwork(
model,
input_infos=input_info_list,
dummy_forward_fn=dummy_forward_fn,
wrap_inputs_fn=wrap_inputs_fn,
wrap_outputs_fn=wrap_outputs_fn,
ignored_scopes=ignored_scopes,
target_scopes=target_scopes,
scopes_without_shape_matching=scopes_without_shape_matching,
original_model_accuracy=original_model_accuracy)
should_init = compression_state is None
builder = create_compression_algorithm_builder(config, should_init)
is_state_loadable = not is_legacy_model_state_dict and compression_state is not None
if is_state_loadable:
builder.load_state(compression_state[BaseController.BUILDER_STATE])
builder.apply_to(compressed_model)
compression_ctrl = builder.build_controller(compressed_model)
if is_state_loadable:
compression_ctrl.load_state(
compression_state[BaseController.CONTROLLER_STATE])
# Required to ensure that the model leaving create_compressed_model has correct compressed graph.
# In particular, this is currently required for correct functioning of RNNs.
compressed_model.rebuild_graph()
try:
if is_legacy_model_state_dict:
from nncf.torch import load_state
state_dict_to_load = compression_state.get('state_dict',
compression_state)
load_state(compressed_model, state_dict_to_load, is_resume=True)
finally:
if dump_graphs and is_main_process():
compressed_model_graph = compressed_model.get_graph()
compressed_model_graph.visualize_graph(
osp.join(config.get("log_dir", "."), "compressed_graph.dot"))
# Synchronize all processes if run in distributed mode
if is_dist_avail_and_initialized():
try:
barrier()
# Exception can be raised during running barrier
# if the backend not in the supported list https://pytorch.org/docs/stable/distributed.html
except RuntimeError as err:
nncf_logger.warning(err)
nncf_logger.warning(
"NNCF continues work, while does not guarantee that "
"the processes will finish model's compression at the same time. "
"If your training pipeline demands the processes be synchronized, please, "
"keep attention to that error")
return compression_ctrl, compressed_model
compressed_model.get_tracing_context().disable_trace_dynamic_graph()
return compression_ctrl, compressed_model