def test_activation_shape_tracing(input_shape: Tuple):
model = ModelForTest()
input_info = ModelInputInfo(input_shape)
graph_builder = GraphBuilder(create_dummy_forward_fn([
input_info,
]))
graph = graph_builder.build_graph(model)
shape1 = (input_shape[0], ModelForTest.CONV1_OUT_CHANNELS, input_shape[2],
input_shape[3])
ref_node_ids_and_output_shapes = [
# TODO: extend with checking input tensor size once proper input node marking is implemented
("0 ModelForTest/Conv2d[conv1]/conv2d", [shape1]),
("1 ModelForTest/BatchNorm2d[bn1]/batch_norm", [shape1]),
("2 ModelForTest/ReLU[relu1]/RELU", [shape1, shape1]),
("3 ModelForTest/max_pool2d",
[(shape1[0], shape1[1], shape1[2] // ModelForTest.MAXPOOL_SIZE,
shape1[3] // ModelForTest.MAXPOOL_SIZE)]),
("4 ModelForTest/ConvTranspose2d[convt1]/conv_transpose2d",
[input_shape]),
("5 ModelForTest/cat",
[(input_shape[0], ModelForTest.CONV2_IN_CHANNELS, input_shape[2],
input_shape[3])])
# TODO: extend with checking output tensor size once proper output node marking is implemented
]
for node_id, ref_output_shapes in ref_node_ids_and_output_shapes:
# pylint:disable=protected-access
output_edges = graph._get_nncf_graph_pattern_input_output([
node_id,
]).output_edges
output_shapes = [x.tensor_shape for x in output_edges]
assert output_shapes == ref_output_shapes, "Failed for {}".format(
node_id)
def _get_dummy_forward_fn_for_graph_building(self, with_input_tracing):
if self._dummy_forward_fn is None:
return create_dummy_forward_fn(
self.input_infos,
with_input_tracing=with_input_tracing,
wrap_inputs_fn=self._wrap_inputs_fn)
return self._dummy_forward_fn
def get_all_node_names(model, input_sample_size, builder=None):
if not builder:
builder = GraphBuilder(
create_dummy_forward_fn([
ModelInputInfo(input_sample_size),
]))
graph = builder.build_graph(model)
return [
node_name.split(' ', 1)[1] for node_name in graph.get_all_node_keys()
]
def test_build_graph(self, desc: ModelDesc):
net = desc.model_builder()
input_sample_sizes = desc.input_sample_sizes
if isinstance(input_sample_sizes, tuple):
input_info_list = [
ModelInputInfo(sample_size)
for sample_size in input_sample_sizes
]
else:
input_info_list = [ModelInputInfo(input_sample_sizes)]
dummy_forward_fn = desc.dummy_forward_fn
if not dummy_forward_fn:
dummy_forward_fn = create_dummy_forward_fn(input_info_list)
graph_builder = GraphBuilder(custom_forward_fn=dummy_forward_fn)
graph = graph_builder.build_graph(net)
check_graph(graph, desc.dot_filename, 'original')
def test_ambiguous_function():
class Model(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList(
[nn.Conv2d(1, 1, 1), nn.Conv2d(1, 1, 1)])
def forward(self, x):
for layer in self.layers:
x = F.relu(layer(x))
mod = Model()
input_info = ModelInputInfo([1, 1, 1, 1])
graph_builder = GraphBuilder(custom_forward_fn=create_dummy_forward_fn([
input_info,
]))
graph = graph_builder.build_graph(mod)
unique_op_exec_contexts = set()
# pylint:disable=protected-access
for _, node in graph._nx_graph.nodes.items():
node_op_exec_context = node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR]
assert node_op_exec_context not in unique_op_exec_contexts
def create_compressed_model(model: Module, config: NNCFConfig,
resuming_state_dict: dict = None,
dummy_forward_fn: Callable[[Module], Any] = 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 resuming_state_dict: A PyTorch state dict object to load (strictly) into the compressed model after
building.
:param dummy_forward_fn: 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.
: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 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.dump_graph(osp.join(config.get("log_dir", "."),
"original_graph.dot"),
extended=True)
if is_debug():
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')
compressed_model = NNCFNetwork(
model,
input_infos=input_info_list,
dummy_forward_fn=dummy_forward_fn,
ignored_scopes=ignored_scopes,
target_scopes=target_scopes,
scopes_without_shape_matching=scopes_without_shape_matching)
should_init = resuming_state_dict is None
compression_algo_builder_list = create_compression_algorithm_builders(
config, should_init=should_init)
for builder in compression_algo_builder_list:
compressed_model = builder.apply_to(compressed_model)
compression_ctrl = compressed_model.commit_compression_changes()
if dump_graphs and is_main_process() and compression_algo_builder_list:
if dummy_forward_fn is None:
compressed_graph_builder = GraphBuilder(
custom_forward_fn=create_dummy_forward_fn(
input_info_list, with_input_tracing=False))
else:
compressed_graph_builder = GraphBuilder(
custom_forward_fn=dummy_forward_fn)
graph = compressed_graph_builder.build_graph(
compressed_model, compressed_model.get_tracing_context())
graph.dump_graph(osp.join(config.get("log_dir", "."),
"compressed_graph.dot"),
extended=True)
if resuming_state_dict is not None:
load_state(compressed_model, resuming_state_dict, is_resume=True)
return compression_ctrl, compressed_model
def create_compressed_model(model: Module, config: NNCFConfig,
resuming_state_dict: dict = None,
dummy_forward_fn: Callable[[Module], Any] = None,
wrap_inputs_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 resuming_state_dict: A PyTorch state dict object to load (strictly) into the compressed model after
building.
: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.
: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.
: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."""
# 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')
compressed_model = NNCFNetwork(
model,
input_infos=input_info_list,
dummy_forward_fn=dummy_forward_fn,
wrap_inputs_fn=wrap_inputs_fn,
ignored_scopes=ignored_scopes,
target_scopes=target_scopes,
scopes_without_shape_matching=scopes_without_shape_matching)
should_init = resuming_state_dict is None
compression_algo_builder_list = create_compression_algorithm_builders(
config, should_init=should_init)
for builder in compression_algo_builder_list:
compressed_model = builder.apply_to(compressed_model)
compression_ctrl = compressed_model.commit_compression_changes()
try:
if resuming_state_dict is not None:
load_state(compressed_model, resuming_state_dict, is_resume=True)
finally:
if dump_graphs and is_main_process() and compression_algo_builder_list:
if dummy_forward_fn is None:
compressed_graph_builder = GraphBuilder(
custom_forward_fn=create_dummy_forward_fn(
input_info_list, with_input_tracing=False))
else:
compressed_graph_builder = GraphBuilder(
custom_forward_fn=dummy_forward_fn)
graph = compressed_graph_builder.build_graph(
compressed_model, compressed_model.get_tracing_context())
graph.visualize_graph(
osp.join(config.get("log_dir", "."), "compressed_graph.dot"))
return compression_ctrl, compressed_model
def test_build_graph(self, model_name, model_builder, input_size):
net = model_builder()
graph_builder = GraphBuilder(create_dummy_forward_fn([ModelInputInfo(input_size), ]))
graph = graph_builder.build_graph(net)
check_graph(graph, model_name, 'original')
