Compression Setting¶
To enhance the compression compatibility of any type of module, NNI 3.0 introduces the notion of setting
,
which is a pre-established template utilized to depict compression information.
The primary objective of this pre-established template is to support shorthand writing in the config_list
.
The list of NNI default supported module types can be accessed via this link.
Please review the registry
in PruningSetting
, QuantizationSetting
and DistillationSetting
to see the supported module type and its default setting.
It should be noted that DistillationSetting
will automatically register a default output setting for all module types,
which implies that distilling any module output is available by design.
Register Setting¶
If you discover that the module type you intend to compress is unavailable, you may register it into the appropriate compression setting. However, if the registered module type is already present in the registry, the new setting will directly replace the old one. In the subsequent sections, you will learn how to register settings in various compression scenarios.
You can find all the supported compression setting keys and their explanations here.
If you only want to make temporary settings for a specific layer without affecting other default templates of the same module type, please refer to section.
Assume a customized module type and we want to compress it:
class CustomizedModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.p1 = torch.nn.Parameter(200, 100)
self.p2 = torch.nn.Parameter(200, 50)
def forward(self, x, y):
return torch.matmul(self.p1, x) + torch.matmul(self.p2, y)
Pruning Setting¶
Here is a default setting for pruning module weight and bias.
default_setting = {
'weight': {
'sparse_ratio': None,
'max_sparse_ratio': None,
'min_sparse_ratio': None,
'sparse_threshold': None,
'global_group_id': None,
'dependency_group_id': None,
'granularity': 'default',
'internal_metric_block': None,
'apply_method': 'mul',
},
'bias': {
'align': {
'module_name': None,
'target_name': 'weight',
'dims': [0],
},
'apply_method': 'mul',
}
}
We can create a setting for CustomizedModule
by following the above default setting.
customized_setting = {
'p1': {
'sparse_ratio': None,
'max_sparse_ratio': None,
'min_sparse_ratio': None,
'sparse_threshold': None,
'global_group_id': None,
'dependency_group_id': None,
'granularity': [1, -1],
'internal_metric_block': None,
'apply_method': 'mul',
},
'p2': {
'sparse_ratio': None,
'max_sparse_ratio': None,
'min_sparse_ratio': None,
'sparse_threshold': None,
'global_group_id': None,
'dependency_group_id': None,
'granularity': [1, -1],
'internal_metric_block': None,
'apply_method': 'mul',
},
'_output_': {
'align': {
'module_name': None,
'target_name': 'p1',
'dims': [0],
},
'apply_method': 'mul',
'granularity': [-1, 1]
}
}
PruningSetting.register('CustomizedModule', customized_setting)
The customized setting means that p1
and p2
will be applied channel-wise masks on the first dim of parameter,
_output_
will be applied channel-wise masks on the second dim of output.
Instead of generating masks by pruning algorithms, the output masks is generated by align with p1
masks on the first dim.
Quantization Setting¶
Here is a default setting for quantizing module inputs, outputs and weight.
default_setting = {
'_input_': {
'quant_dtype': None,
'quant_scheme': None,
'granularity': 'default',
'apply_method': 'clamp_round',
},
'weight': {
'quant_dtype': None,
'quant_scheme': None,
'granularity': 'default',
'apply_method': 'clamp_round',
},
'_output_': {
'quant_dtype': None,
'quant_scheme': None,
'granularity': 'default',
'apply_method': 'clamp_round',
}
}
Just modified the keys and registered it to QuantizationSetting
are all you need for quantizing this module.
customized_setting = {
'_input_': {
'quant_dtype': None,
'quant_scheme': None,
'granularity': [-1, 1],
'apply_method': 'clamp_round',
},
'p1': {
'quant_dtype': None,
'quant_scheme': None,
'granularity': [1, -1],
'apply_method': 'clamp_round',
},
'p2': {
'quant_dtype': None,
'quant_scheme': None,
'granularity': [1, -1],
'apply_method': 'clamp_round',
},
'_output_': {
'quant_dtype': None,
'quant_scheme': None,
'granularity': [-1, 1],
'apply_method': 'clamp_round',
}
}
QuantizationSetting.register('CustomizedModule', customized_setting)
Temporarily Setting Update¶
Sometimes we just want to temporarily modify the setting template and don't want to make a global change.
Then we could directly write full setting in config_list
to achieve this.
For example, if the compressed model has conv-bn-relu pattern, and for a better pruning simulation and performance, we want to mask the batchnorm on the channels convolution masked. Then we could temporarily make batchnorm masks align with convolution layer weight masks.
config_list = [{
'op_names': ['conv1', 'conv2'],
'granularity': 'out_channel',
}, {
'op_names': ['bn1'],
'target_settings': {
'weight': {
'align': {
'module_name': 'conv1',
'target_name': 'weight',
'dims': [0],
},
'granularity': 'out_channel',
}
}
}, {
'op_names': ['bn2'],
'target_settings': {
'weight': {
'align': {
'module_name': 'conv2',
'target_name': 'weight',
'dims': [0],
},
'granularity': 'out_channel',
}
}
}]