Skip to content

GitLab

Unverified Commit 86335921 authored by lin bin's avatar lin bin Committed by GitHub
Browse files

[Model Compression Quantization] Unify variable name (#3990)

parent e5c3ac63
Hide whitespace changes
Inline Side-by-side
Showing with 176 additions and 166 deletions
examples/model_compress/quantization/mixed_precision_speedup_mnist.py
View file @ 86335921
......@@ -58,10 +58,10 @@ def post_training_quantization_example(train_loader, test_loader, device):
model = NaiveModel()
config = {
'conv1':{'weight_bit':8, 'activation_bit':8},
'conv2':{'weight_bit':32, 'activation_bit':32},
'fc1':{'weight_bit':16, 'activation_bit':16},
'fc2':{'weight_bit':8, 'activation_bit':8}
'conv1':{'weight_bits':8, 'output_bits':8},
'conv2':{'weight_bits':32, 'output_bits':32},
'fc1':{'weight_bits':16, 'output_bits':16},
'fc2':{'weight_bits':8, 'output_bits':8}
}
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
......@@ -102,8 +102,10 @@ def quantization_aware_training_example(train_loader, test_loader, device):
]
# finetune the model by using QAT
# enable batchnorm folding mode
dummy_input = torch.randn(1, 1, 28, 28)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
quantizer = QAT_Quantizer(model, configure_list, optimizer)
quantizer = QAT_Quantizer(model, configure_list, optimizer, dummy_input=dummy_input)
quantizer.compress()
model.to(device)
......
nni/algorithms/compression/pytorch/quantization/quantizers.py
View file @ 86335921
......@@ -124,7 +124,7 @@ class QATGrad(QuantGrad):
class ObserverQuantizer(Quantizer):
"""This quantizer uses observers to record weight/activation statistics to get quantization information.
"""This quantizer uses observers to record weight/output statistics to get quantization information.
The whole process can be divided into three steps:
1. It will register observers to the place where quantization would happen (just like registering hooks).
2. The observers would record tensors' statistics during calibration.
......@@ -140,7 +140,7 @@ class ObserverQuantizer(Quantizer):
# TODO:
# 1. support dtype and qscheme customization through config_list. Current settings:
# weight observer : per_tensor_symmetric, qint8
# activation observer : per_tensor_affine, quint8, reduce_range=True
# output observer : per_tensor_affine, quint8, reduce_range=True
# 2. add more kinds of observers, such as Kullback-Leibler divergence.
# 3. add batch normalization folding
assert not model.training, "Currently the observer quantizer only works in evaluation mode."
......@@ -148,8 +148,8 @@ class ObserverQuantizer(Quantizer):
self.device = next(model.parameters()).device
modules_to_compress = self.get_modules_to_compress()
all_observers = defaultdict(dict)
weight_q_min, weight_q_max = -127, 127
activation_q_min, activation_q_max = 0, 127 # reduce_range is set to True
weight_qmin, weight_qmax = -127, 127
output_qmin, output_qmax = 0, 127 # reduce_range is set to True
self.compressed = False
for layer, config in modules_to_compress:
......@@ -157,16 +157,16 @@ class ObserverQuantizer(Quantizer):
module = layer.module
if "weight" in config.get("quant_types", []):
all_observers[layer_name]["weight"] = default_weight_observer()
setattr(module, "weight_qmax", weight_q_max)
setattr(module, "weight_qmin", weight_q_min)
setattr(module, "weight_qmax", weight_qmax)
setattr(module, "weight_qmin", weight_qmin)
if "input" in config.get("quant_types", []):
all_observers[layer_name]["input"] = default_histogram_observer()
setattr(module, "input_qmax", activation_q_max)
setattr(module, "input_qmin", activation_q_min)
setattr(module, "input_qmax", output_qmax)
setattr(module, "input_qmin", output_qmin)
if "output" in config.get("quant_types", []):
all_observers[layer_name]["output"] = default_histogram_observer()
setattr(module, "output_qmax", activation_q_max)
setattr(module, "output_qmin", activation_q_min)
setattr(module, "output_qmax", output_qmax)
setattr(module, "output_qmin", output_qmin)
self.all_observers = all_observers
self.bound_model.to(self.device)
......@@ -306,29 +306,29 @@ class ObserverQuantizer(Quantizer):
if hasattr(module, 'weight_scale') or hasattr(module, 'input_scale') or hasattr(module, 'output_scale'):
calibration_config[name] = {}
if hasattr(module, 'weight_scale'):
calibration_config[name]['weight_bit'] = 8
calibration_config[name]['weight_bits'] = 8
val = float(module.weight_scale * module.weight_qmax)
calibration_config[name]['tracked_max_weight'] = val
calibration_config[name]['tracked_min_weight'] = -val
calibration_config[name]['tracked_weight_qmin'] = -127
calibration_config[name]['tracked_weight_qmax'] = 127
calibration_config[name]['tracked_qmin_weight'] = -127
calibration_config[name]['tracked_qmax_weight'] = 127
# refactor these magic numbers when customizations of dtype and qscheme are ready.
if hasattr(module, 'input_scale'):
calibration_config[name]['input_bit'] = 8
calibration_config[name]['input_bits'] = 8
max_input = float(module.input_scale * (module.input_qmax - module.input_zero_point))
min_input = float(module.input_scale * (module.input_qmin - module.input_zero_point))
calibration_config[name]['tracked_min_input'] = min_input
calibration_config[name]['tracked_max_input'] = max_input
calibration_config[name]['tracked_input_qmin'] = 0
calibration_config[name]['tracked_input_qmax'] = 127
calibration_config[name]['tracked_qmin_input'] = 0
calibration_config[name]['tracked_qmax_input'] = 127
if hasattr(module, 'output_scale'):
calibration_config[name]['activation_bit'] = 8
calibration_config[name]['output_bits'] = 8
max_input = float(module.output_scale * (module.output_qmax - module.output_zero_point))
min_input = float(module.output_scale * (module.output_qmin - module.output_zero_point))
calibration_config[name]['tracked_min_activation'] = min_input
calibration_config[name]['tracked_max_activation'] = max_input
calibration_config[name]['tracked_activation_qmin'] = 0
calibration_config[name]['tracked_activation_qmax'] = 127
calibration_config[name]['tracked_min_output'] = min_input
calibration_config[name]['tracked_max_output'] = max_input
calibration_config[name]['tracked_qmin_output'] = 0
calibration_config[name]['tracked_qmax_output'] = 127
self._del_simulated_attr(module)
self.export_model_save(self.bound_model, model_path, calibration_config, calibration_path, onnx_path,
......@@ -354,7 +354,7 @@ class QAT_Quantizer(Quantizer):
http://openaccess.thecvf.com/content_cvpr_2018/papers/Jacob_Quantization_and_Training_CVPR_2018_paper.pdf
"""
def __init__(self, model, config_list, optimizer=None, dummy_input=None):
def __init__(self, model, config_list, optimizer, dummy_input=None):
"""
Parameters
----------
......@@ -370,7 +370,7 @@ class QAT_Quantizer(Quantizer):
when the type is int, all quantization types share same bits length
- quant_start_step : int
disable quantization until model are run by certain number of steps, this allows the network to enter a more stable
state where activation quantization ranges do not exclude a significant fraction of values, default value is 0
state where output quantization ranges do not exclude a significant fraction of values, default value is 0
- op_types : list of string
types of nn.module you want to apply quantization, eg. 'Conv2d'
- dummy_input : tuple of tensor
......@@ -379,6 +379,7 @@ class QAT_Quantizer(Quantizer):
given, the batch normalization folding would be disabled.
"""
assert isinstance(optimizer, torch.optim.Optimizer), "unrecognized optimizer type"
super().__init__(model, config_list, optimizer, dummy_input)
self.quant_grad = QATGrad.apply
modules_to_compress = self.get_modules_to_compress()
......@@ -389,22 +390,22 @@ class QAT_Quantizer(Quantizer):
layer.module.register_buffer("scale", torch.Tensor([1.0]))
layer.module.register_buffer('ema_decay', torch.Tensor([0.99]))
if "weight" in config.get("quant_types", []):
layer.module.register_buffer('weight_bit', torch.zeros(1))
layer.module.register_buffer('weight_bits', torch.zeros(1))
layer.module.register_buffer('tracked_min_input', torch.zeros(1))
layer.module.register_buffer('tracked_max_input', torch.zeros(1))
if "output" in config.get("quant_types", []):
layer.module.register_buffer('activation_bit', torch.zeros(1))
layer.module.register_buffer('tracked_min_activation', torch.zeros(1))
layer.module.register_buffer('tracked_max_activation', torch.zeros(1))
layer.module.register_buffer('output_bits', torch.zeros(1))
layer.module.register_buffer('tracked_min_output', torch.zeros(1))
layer.module.register_buffer('tracked_max_output', torch.zeros(1))
self.bound_model.to(device)
def _del_simulated_attr(self, module):
"""
delete redundant parameters in quantize module
"""
del_attr_list = ['old_weight', 'old_bias', 'ema_decay', 'tracked_min_activation', 'tracked_max_activation',
'tracked_min_input', 'tracked_max_input', 'scale', 'zero_point', 'weight_bit',
'activation_bit', 'BN_FOLD_TAG']
del_attr_list = ['old_weight', 'old_bias', 'ema_decay', 'tracked_min_output', 'tracked_max_output',
'tracked_min_input', 'tracked_max_input', 'scale', 'zero_point', 'weight_bits',
'output_bits', 'BN_FOLD_TAG']
for attr in del_attr_list:
if hasattr(module, attr):
delattr(module, attr)
......@@ -506,7 +507,7 @@ class QAT_Quantizer(Quantizer):
module.scale, module.zero_point = update_quantization_param(weight_bits, rmin, rmax)
weight = self._quantize(weight_bits, module, weight)
weight = self._dequantize(module, weight)
module.weight_bit = torch.Tensor([weight_bits])
module.weight_bits = torch.Tensor([weight_bits])
wrapper.module.weight = weight
return weight
......@@ -514,23 +515,23 @@ class QAT_Quantizer(Quantizer):
config = wrapper.config
module = wrapper.module
output_bits = get_bits_length(config, 'output')
module.activation_bit = torch.Tensor([output_bits])
module.output_bits = torch.Tensor([output_bits])
quant_start_step = config.get('quant_start_step', 0)
assert output_bits >= 1, "quant bits length should be at least 1"
if quant_start_step > self.bound_model.steps:
module.tracked_min_activation, module.tracked_max_activation = torch.min(output), torch.max(output)
module.tracked_min_output, module.tracked_max_output = torch.min(output), torch.max(output)
return output
# we dont update output quantization parameters in evaluation stage
if wrapper.training:
current_min, current_max = torch.min(output), torch.max(output)
module.tracked_min_activation = update_ema(module.tracked_min_activation, current_min,
module.tracked_min_output = update_ema(module.tracked_min_output, current_min,
module.ema_decay)
module.tracked_max_activation = update_ema(module.tracked_max_activation, current_max,
module.tracked_max_output = update_ema(module.tracked_max_output, current_max,
module.ema_decay)
module.scale, module.zero_point = update_quantization_param(
output_bits, module.tracked_min_activation, module.tracked_max_activation)
output_bits, module.tracked_min_output, module.tracked_max_output)
out = self._quantize(output_bits, module, output)
out = self._dequantize(module, out)
return out
......@@ -562,10 +563,10 @@ class QAT_Quantizer(Quantizer):
calibration_config = {}
for name, module in self.bound_model.named_modules():
if hasattr(module, 'weight_bit') or hasattr(module, 'activation_bit'):
if hasattr(module, 'weight_bits') or hasattr(module, 'output_bits'):
calibration_config[name] = {}
if hasattr(module, 'weight_bit'):
calibration_config[name]['weight_bit'] = int(module.weight_bit)
if hasattr(module, 'weight_bits'):
calibration_config[name]['weight_bits'] = int(module.weight_bits)
calibration_config[name]['tracked_min_input'] = float(module.tracked_min_input)
calibration_config[name]['tracked_max_input'] = float(module.tracked_max_input)
......@@ -585,10 +586,10 @@ class QAT_Quantizer(Quantizer):
else:
setattr(module, 'bias', None)
if hasattr(module, 'activation_bit'):
calibration_config[name]['activation_bit'] = int(module.activation_bit)
calibration_config[name]['tracked_min_activation'] = float(module.tracked_min_activation)
calibration_config[name]['tracked_max_activation'] = float(module.tracked_max_activation)
if hasattr(module, 'output_bits'):
calibration_config[name]['output_bits'] = int(module.output_bits)
calibration_config[name]['tracked_min_output'] = float(module.tracked_min_output)
calibration_config[name]['tracked_max_output'] = float(module.tracked_max_output)
self._del_simulated_attr(module)
self.export_model_save(self.bound_model, model_path, calibration_config, calibration_path, onnx_path, input_shape, device)
......@@ -642,20 +643,21 @@ class DoReFaQuantizer(Quantizer):
(https://arxiv.org/abs/1606.06160)
"""
def __init__(self, model, config_list, optimizer=None):
def __init__(self, model, config_list, optimizer):
assert isinstance(optimizer, torch.optim.Optimizer), "unrecognized optimizer type"
super().__init__(model, config_list, optimizer)
device = next(model.parameters()).device
modules_to_compress = self.get_modules_to_compress()
for layer, config in modules_to_compress:
if "weight" in config.get("quant_types", []):
layer.module.register_buffer('weight_bit', torch.zeros(1))
layer.module.register_buffer('weight_bits', torch.zeros(1))
self.bound_model.to(device)
def _del_simulated_attr(self, module):
"""
delete redundant parameters in quantize module
"""
del_attr_list = ['old_weight', 'weight_bit']
del_attr_list = ['old_weight', 'weight_bits']
for attr in del_attr_list:
if hasattr(module, attr):
delattr(module, attr)
......@@ -689,7 +691,7 @@ class DoReFaQuantizer(Quantizer):
weight = self.quantize(weight, weight_bits)
weight = 2 * weight - 1
wrapper.module.weight = weight
wrapper.module.weight_bit = torch.Tensor([weight_bits])
wrapper.module.weight_bits = torch.Tensor([weight_bits])
# wrapper.module.weight.data = weight
return weight
......@@ -725,9 +727,9 @@ class DoReFaQuantizer(Quantizer):
calibration_config = {}
for name, module in self.bound_model.named_modules():
if hasattr(module, 'weight_bit'):
if hasattr(module, 'weight_bits'):
calibration_config[name] = {}
calibration_config[name]['weight_bit'] = int(module.weight_bit)
calibration_config[name]['weight_bits'] = int(module.weight_bits)
self._del_simulated_attr(module)
self.export_model_save(self.bound_model, model_path, calibration_config, calibration_path, onnx_path, input_shape, device)
......@@ -745,25 +747,26 @@ class ClipGrad(QuantGrad):
class BNNQuantizer(Quantizer):
"""Binarized Neural Networks, as defined in:
Binarized Neural Networks: Training Deep Neural Networks with Weights and Activations Constrained to +1 or -1
Binarized Neural Networks: Training Deep Neural Networks with Weights and Outputs Constrained to +1 or -1
(https://arxiv.org/abs/1602.02830)
"""
def __init__(self, model, config_list, optimizer=None):
def __init__(self, model, config_list, optimizer):
assert isinstance(optimizer, torch.optim.Optimizer), "unrecognized optimizer type"
super().__init__(model, config_list, optimizer)
device = next(model.parameters()).device
self.quant_grad = ClipGrad.apply
modules_to_compress = self.get_modules_to_compress()
for layer, config in modules_to_compress:
if "weight" in config.get("quant_types", []):
layer.module.register_buffer('weight_bit', torch.zeros(1))
layer.module.register_buffer('weight_bits', torch.zeros(1))
self.bound_model.to(device)
def _del_simulated_attr(self, module):
"""
delete redundant parameters in quantize module
"""
del_attr_list = ['old_weight', 'weight_bit']
del_attr_list = ['old_weight', 'weight_bits']
for attr in del_attr_list:
if hasattr(module, attr):
delattr(module, attr)
......@@ -796,7 +799,7 @@ class BNNQuantizer(Quantizer):
# remove zeros
weight[weight == 0] = 1
wrapper.module.weight = weight
wrapper.module.weight_bit = torch.Tensor([1.0])
wrapper.module.weight_bits = torch.Tensor([1.0])
return weight
def quantize_output(self, output, wrapper, **kwargs):
......@@ -832,9 +835,9 @@ class BNNQuantizer(Quantizer):
calibration_config = {}
for name, module in self.bound_model.named_modules():
if hasattr(module, 'weight_bit'):
if hasattr(module, 'weight_bits'):
calibration_config[name] = {}
calibration_config[name]['weight_bit'] = int(module.weight_bit)
calibration_config[name]['weight_bits'] = int(module.weight_bits)
self._del_simulated_attr(module)
self.export_model_save(self.bound_model, model_path, calibration_config, calibration_path, onnx_path, input_shape, device)
......@@ -848,7 +851,7 @@ class LsqQuantizer(Quantizer):
https://arxiv.org/pdf/1902.08153.pdf
"""
def __init__(self, model, config_list, optimizer=None):
def __init__(self, model, config_list, optimizer):
"""
Parameters
----------
......@@ -864,10 +867,11 @@ class LsqQuantizer(Quantizer):
when the type is int, all quantization types share same bits length
- quant_start_step : int
disable quantization until model are run by certain number of steps, this allows the network to enter a more stable
state where activation quantization ranges do not exclude a significant fraction of values, default value is 0
state where output quantization ranges do not exclude a significant fraction of values, default value is 0
- op_types : list of string
types of nn.module you want to apply quantization, eg. 'Conv2d'
"""
assert isinstance(optimizer, torch.optim.Optimizer), "unrecognized optimizer type"
super().__init__(model, config_list, optimizer)
device = next(model.parameters()).device
self.quant_grad = QuantForward()
......@@ -877,10 +881,10 @@ class LsqQuantizer(Quantizer):
if "weight" in config.get("quant_types", []):
layer.module.register_parameter("weight_scale", torch.nn.Parameter(torch.Tensor([1.0])))
# todo: support per-channel quantization for weight since TensorRT use it for conv weight
q_bit = get_bits_length(config, "weight")
layer.module.register_buffer('weight_bit', torch.Tensor([q_bit]))
qmax = 2 ** (q_bit - 1) - 1
qmin = -2 ** (q_bit - 1)
q_bits = get_bits_length(config, "weight")
layer.module.register_buffer('weight_bits', torch.Tensor([q_bits]))
qmax = 2 ** (q_bits - 1) - 1
qmin = -2 ** (q_bits - 1)
init_weight_scale = layer.module.weight.data.detach().abs().mean() * 2 / (qmax ** 0.5)
layer.module.weight_scale = torch.nn.Parameter(init_weight_scale)
layer.module.weight_qmax = qmax
......@@ -889,12 +893,12 @@ class LsqQuantizer(Quantizer):
self.optimizer.add_param_group({"params": layer.module.weight_scale})
if "output" in config.get("quant_types", []):
# scale of activation will be initialized using the first batch data
# scale of output will be initialized using the first batch data
layer.module.register_parameter("output_scale", torch.nn.Parameter(torch.Tensor([1.0])))
q_bit = get_bits_length(config, "output")
layer.module.register_buffer('output_bit', torch.Tensor([q_bit]))
qmax = 2 ** (q_bit - 1) - 1
qmin = -2 ** (q_bit - 1)
q_bits = get_bits_length(config, "output")
layer.module.register_buffer('output_bits', torch.Tensor([q_bits]))
qmax = 2 ** (q_bits - 1) - 1
qmin = -2 ** (q_bits - 1)
layer.module.output_qmax = qmax
layer.module.output_qmin = qmin
......@@ -903,10 +907,10 @@ class LsqQuantizer(Quantizer):
if "input" in config.get("quant_types", []):
# scale of input will be initialized using the first batch data
layer.module.register_parameter("input_scale", torch.nn.Parameter(torch.Tensor([1.0])))
q_bit = get_bits_length(config, "input")
layer.module.register_buffer('input_bit', torch.Tensor([q_bit]))
qmax = 2 ** (q_bit - 1) - 1
qmin = -2 ** (q_bit - 1)
q_bits = get_bits_length(config, "input")
layer.module.register_buffer('input_bits', torch.Tensor([q_bits]))
qmax = 2 ** (q_bits - 1) - 1
qmin = -2 ** (q_bits - 1)
layer.module.input_qmax = qmax
layer.module.input_qmin = qmin
......@@ -1011,18 +1015,18 @@ class LsqQuantizer(Quantizer):
calibration_config = {}
for name, module in self.bound_model.named_modules():
if hasattr(module, 'input_bit') or hasattr(module, 'output_bit'):
if hasattr(module, 'input_bits') or hasattr(module, 'output_bits'):
calibration_config[name] = {}
if hasattr(module, 'weight_bit'):
calibration_config[name]['weight_bit'] = int(module.weight_bit)
if hasattr(module, 'weight_bits'):
calibration_config[name]['weight_bits'] = int(module.weight_bits)
abs_max_input = float(module.input_scale * module.input_qmax)
calibration_config[name]['tracked_min_input'] = -abs_max_input
calibration_config[name]['tracked_max_input'] = abs_max_input
if hasattr(module, 'output_bit'):
calibration_config[name]['activation_bit'] = int(module.output_bit)
if hasattr(module, 'output_bits'):
calibration_config[name]['output_bits'] = int(module.output_bits)
abs_max_output = float(module.output_scale * module.output_qmax)
calibration_config[name]['tracked_min_activation'] = -abs_max_output
calibration_config[name]['tracked_max_activation'] = abs_max_output
calibration_config[name]['tracked_min_output'] = -abs_max_output
calibration_config[name]['tracked_max_output'] = abs_max_output
self._del_simulated_attr(module)
self.export_model_save(self.bound_model, model_path, calibration_config, calibration_path, onnx_path,
......@@ -1034,8 +1038,8 @@ class LsqQuantizer(Quantizer):
"""
delete redundant parameters in quantize module
"""
del_attr_list = ['old_weight', 'tracked_min_input', 'tracked_max_input', 'tracked_min_activation', \
'tracked_max_activation', 'output_scale', 'input_scale', 'weight_scale','weight_bit', 'output_bit', 'input_bit']
del_attr_list = ['old_weight', 'tracked_min_input', 'tracked_max_input', 'tracked_min_output', \
'tracked_max_output', 'output_scale', 'input_scale', 'weight_scale','weight_bits', 'output_bits', 'input_bits']
for attr in del_attr_list:
if hasattr(module, attr):
delattr(module, attr)
......
nni/compression/pytorch/compressor.py
View file @ 86335921
......@@ -834,7 +834,7 @@ class QuantGrad(torch.autograd.Function):
@classmethod
def get_bits_length(cls, config, quant_type):
"""
Get bit for quantize config
Get bits for quantize config
Parameters
----------
config : Dict
......
nni/compression/pytorch/quantization_speedup/frontend_to_onnx.py
View file @ 86335921
......@@ -9,26 +9,26 @@ The main function of this page is to convert pytorch model to onnx model.
Convertion from pytorch model to onnx model is primary so that a critical
problem is caused that Layer name of pytorch model fail to convert to onnx
layer name directly. To solve it, we wrap pytorch model in new wrapper which
multiply bit number and input before computation of each op. Only in this
way can onnx model get bit number of corresponded layer.
multiply bits number and input before computation of each op. Only in this
way can onnx model get bits number of corresponded layer.
"""
class LayernameModuleWrapper(torch.nn.Module):
def __init__(self, module, module_bit) -> None:
def __init__(self, module, module_bits) -> None:
"""
Parameters
----------
module : torch.nn.Module
Layer module of pytorch model
module_bit : int
Bit width setting for module
module_bits : int
Bits width setting for module
"""
super().__init__()
self.module = module
self.module_bit = module_bit
self.module_bits = module_bits
def forward(self, inputs):
inputs = inputs*self.module_bit
inputs = inputs*self.module_bits
inputs = self.module(inputs)
return inputs
......@@ -93,14 +93,14 @@ def unwrapper(model_onnx, index2name, config):
def torch_to_onnx(model, config, input_shape, model_path, input_names, output_names):
"""
Convert torch model to onnx model and get layer bit config of onnx model.
Convert torch model to onnx model and get layer bits config of onnx model.
Parameters
----------
model : pytorch model
The model to speed up by quantization
config : dict
Config recording bit number and name of layers
Config recording bits number and name of layers
input_shape : tuple
The input shape of model, shall pass it to torch.onnx.export
model_path : str
......@@ -119,7 +119,7 @@ def torch_to_onnx(model, config, input_shape, model_path, input_names, output_na
"""
# Support Gemm, Conv, Relu, Clip(Relu6) and MaxPool
support_op = [torch.nn.Conv2d, torch.nn.Linear, torch.nn.ReLU, torch.nn.ReLU6, torch.nn.MaxPool2d]
# Transfer bit number to onnx layer by using wrapper
# Transfer bits number to onnx layer by using wrapper
index2name = {}
name2index = {}
if config is not None:
......
nni/compression/pytorch/quantization_speedup/integrated_tensorrt.py
View file @ 86335921
......@@ -31,18 +31,18 @@ Precision_Dict = {
def valid_config(config=None):
"""
This function validates the bit setting configuration
This function validates the bits setting configuration
"""
if config is None:
return
support_bit = [8, 16, 32]
support_bits = [8, 16, 32]
for name in config.keys():
if 'weight_bit' in config[name]:
w_bit = config[name]['weight_bit']
assert w_bit in support_bit, "weight bit should be 8, 16, 32"
if 'activation_bit' in config[name]:
a_bit = config[name]['activation_bit']
assert a_bit in support_bit, "activation bit should be 8, 16, 32"
if 'weight_bits' in config[name]:
w_bits = config[name]['weight_bits']
assert w_bits in support_bits, "weight bits should be 8, 16, 32"
if 'output_bits' in config[name]:
a_bits = config[name]['output_bits']
assert a_bits in support_bits, "output bits should be 8, 16, 32"
def handle_gemm(network, layer_idx, config):
"""
......@@ -55,26 +55,26 @@ def handle_gemm(network, layer_idx, config):
layer_idx : int
layer index of gemm
config : dict
Config recording bit number and name of layers
Config recording bits number and name of layers
"""
layer = network.get_layer(layer_idx)
pre_layer = network.get_layer(layer_idx-1)
next_layer = network.get_layer(layer_idx+1)
# if weight bit exists, set three layers' precision,
# if weight bits exists, set three layers' precision,
# input tensor range and the first two layers' output type
if 'weight_bit' in config[layer.name]:
if 'weight_bits' in config[layer.name]:
assert 'tracked_min_input' in config[layer.name]
assert 'tracked_max_input' in config[layer.name]
w_bit = config[layer.name]['weight_bit']
w_bits = config[layer.name]['weight_bits']
tracked_min_input = config[layer.name]['tracked_min_input']
tracked_max_input = config[layer.name]['tracked_max_input']
# set three layers the same precision
layer.precision = Precision_Dict[w_bit]
pre_layer.precision = Precision_Dict[w_bit]
next_layer.precision = Precision_Dict[w_bit]
layer.precision = Precision_Dict[w_bits]
pre_layer.precision = Precision_Dict[w_bits]
next_layer.precision = Precision_Dict[w_bits]
# set the first two layers' output type
pre_layer.set_output_type(0, Precision_Dict[w_bit])
layer.set_output_type(0, Precision_Dict[w_bit])
pre_layer.set_output_type(0, Precision_Dict[w_bits])
layer.set_output_type(0, Precision_Dict[w_bits])
pre_in_tensor = pre_layer.get_input(0)
in_tensor = layer.get_input(0)
next_in_tensor = next_layer.get_input(0)
......@@ -83,20 +83,20 @@ def handle_gemm(network, layer_idx, config):
in_tensor.dynamic_range = (tracked_min_input, tracked_max_input)
next_in_tensor.dynamic_range = (tracked_min_input, tracked_max_input)
# if activation bit exists, set the last layer's output type output tensor range
if 'activation_bit' in config[layer.name]:
assert 'tracked_min_activation' in config[layer.name]
assert 'tracked_max_activation' in config[layer.name]
a_bit = config[layer.name]['activation_bit']
tracked_min_activation = config[layer.name]['tracked_min_activation']
tracked_max_activation = config[layer.name]['tracked_max_activation']
# if output bits exists, set the last layer's output type output tensor range
if 'output_bits' in config[layer.name]:
assert 'tracked_min_output' in config[layer.name]
assert 'tracked_max_output' in config[layer.name]
a_bits = config[layer.name]['output_bits']
tracked_min_output = config[layer.name]['tracked_min_output']
tracked_max_output = config[layer.name]['tracked_max_output']
# set the last layer's output type
next_layer.set_output_type(0, Precision_Dict[a_bit])
next_layer.set_output_type(0, Precision_Dict[a_bits])
next_out_tensor = next_layer.get_output(0)
# set the last layer's output tensor range
next_out_tensor.dynamic_range = (tracked_min_activation, tracked_max_activation)
next_out_tensor.dynamic_range = (tracked_min_output, tracked_max_output)
def build_engine(model_file, config=None, extra_layer_bit=32, strict_datatype=False, calib=None):
def build_engine(model_file, config=None, extra_layer_bits=32, strict_datatype=False, calib=None):
"""
This function builds an engine from an onnx model with calibration process.
......@@ -105,12 +105,12 @@ def build_engine(model_file, config=None, extra_layer_bit=32, strict_datatype=Fa
model_file : str
The path of onnx model
config : dict
Config recording bit number and name of layers
extra_layer_bit : int
Other layers which are not in config will be quantized to corresponding bit number
Config recording bits number and name of layers
extra_layer_bits : int
Other layers which are not in config will be quantized to corresponding bits number
strict_datatype : bool
Whether constrain layer bit to the number given in config or not. If true, all the layer
will be set to given bit strictly. Otherwise, these layers will be set automatically by
Whether constrain layer bits to the number given in config or not. If true, all the layer
will be set to given bits strictly. Otherwise, these layers will be set automatically by
tensorrt
calib : numpy array
The data using to calibrate quantization model
......@@ -135,14 +135,14 @@ def build_engine(model_file, config=None, extra_layer_bit=32, strict_datatype=Fa
else:
builder.max_workspace_size = common.GiB(4)
if extra_layer_bit == 32 and config is None:
if extra_layer_bits == 32 and config is None:
pass
elif extra_layer_bit == 16 and config is None:
elif extra_layer_bits == 16 and config is None:
if trt_version == TRT8:
trt_config.set_flag(trt.BuilderFlag.FP16)
else:
builder.fp16_mode = True
elif extra_layer_bit == 8 and config is None:
elif extra_layer_bits == 8 and config is None:
# entire model in 8bit mode
if trt_version == TRT8:
trt_config.set_flag(trt.BuilderFlag.INT8)
......@@ -180,15 +180,15 @@ def build_engine(model_file, config=None, extra_layer_bit=32, strict_datatype=Fa
break
layer = network.get_layer(i)
if layer.name in config:
w_bit = config[layer.name]['weight_bit']
a_bit = config[layer.name]['activation_bit']
layer.precision = Precision_Dict[w_bit]
layer.set_output_type(0, Precision_Dict[a_bit])
w_bits = config[layer.name]['weight_bits']
a_bits = config[layer.name]['output_bits']
layer.precision = Precision_Dict[w_bits]
layer.set_output_type(0, Precision_Dict[a_bits])
else:
# This implementation may be incorrect when output number > 1
for i in range(network.num_layers):
if config is None:
# no low bit layer need to be set, keep original model
# no low bits layer need to be set, keep original model
break
layer = network.get_layer(i)
if layer.name not in config:
......@@ -198,37 +198,37 @@ def build_engine(model_file, config=None, extra_layer_bit=32, strict_datatype=Fa
handle_gemm(network, i, config)
continue
# If weight_bit exists in config, set layer precision and layer's input tensor dynamic range.
if 'weight_bit' in config[layer.name]:
# If weight_bits exists in config, set layer precision and layer's input tensor dynamic range.
if 'weight_bits' in config[layer.name]:
assert 'tracked_min_input' in config[layer.name]
assert 'tracked_max_input' in config[layer.name]
w_bit = config[layer.name]['weight_bit']
w_bits = config[layer.name]['weight_bits']
tracked_min_input = config[layer.name]['tracked_min_input']
tracked_max_input = config[layer.name]['tracked_max_input']
layer.precision = Precision_Dict[w_bit]
layer.precision = Precision_Dict[w_bits]
in_tensor = layer.get_input(0)
in_tensor.dynamic_range = (tracked_min_input, tracked_max_input)
# If activation exists in config, set layer output type and layer's output tensor dynamic range.
if 'activation_bit' in config[layer.name]:
assert 'tracked_min_activation' in config[layer.name]
assert 'tracked_max_activation' in config[layer.name]
a_bit = config[layer.name]['activation_bit']
tracked_min_activation = config[layer.name]['tracked_min_activation']
tracked_max_activation = config[layer.name]['tracked_max_activation']
layer.set_output_type(0, Precision_Dict[a_bit])
# If output exists in config, set layer output type and layer's output tensor dynamic range.
if 'output_bits' in config[layer.name]:
assert 'tracked_min_output' in config[layer.name]
assert 'tracked_max_output' in config[layer.name]
a_bits = config[layer.name]['output_bits']
tracked_min_output = config[layer.name]['tracked_min_output']
tracked_max_output = config[layer.name]['tracked_max_output']
layer.set_output_type(0, Precision_Dict[a_bits])
out_tensor = layer.get_output(0)
out_tensor.dynamic_range = (tracked_min_activation, tracked_max_activation)
out_tensor.dynamic_range = (tracked_min_output, tracked_max_output)
# Build engine and do int8 calibration.
if trt_version == TRT8:
engine = builder.build_engine(network, trt_config)
else:
engine.builder.build_cuda_engine(network)
engine = builder.build_cuda_engine(network)
return engine
class ModelSpeedupTensorRT(BaseModelSpeedup):
def __init__(self, model, input_shape, config=None, onnx_path="default_model.onnx", extra_layer_bit=32, strict_datatype=True,
def __init__(self, model, input_shape, config=None, onnx_path="default_model.onnx", extra_layer_bits=32, strict_datatype=True,
calibrate_type=CalibrateType.ENTROPY2, calib_data_loader=None, calibration_cache = "calibration.cache", batchsize=1,
input_names=["actual_input_1"], output_names=["output1"]):
"""
......@@ -239,14 +239,14 @@ class ModelSpeedupTensorRT(BaseModelSpeedup):
input_shape : tuple
The input shape of model, shall pass it to torch.onnx.export.
config : dict
Config recording bit number and name of layers.
Config recording bits number and name of layers.
onnx_path : str
The path user want to store onnx model which is converted from pytorch model.
extra_layer_bit : int
Other layers which are not in config will be quantized to corresponding bit number.
extra_layer_bits : int
Other layers which are not in config will be quantized to corresponding bits number.
strict_datatype : bool
Whether constrain layer bit to the number given in config or not. If true, all the layer
will be set to given bit strictly. Otherwise, these layers will be set automatically by
Whether constrain layer bits to the number given in config or not. If true, all the layer
will be set to given bits strictly. Otherwise, these layers will be set automatically by
tensorrt.
calibrate_type : tensorrt.tensorrt.CalibrationAlgoType
The algorithm of calibrating. Please refer to https://docs.nvidia.com/deeplearning/
......@@ -267,7 +267,7 @@ class ModelSpeedupTensorRT(BaseModelSpeedup):
self.onnx_path = onnx_path
self.input_shape = input_shape
self.config = config
self.extra_layer_bit = extra_layer_bit
self.extra_layer_bits = extra_layer_bits
self.strict_datatype = strict_datatype
self.calibrate_type = calibrate_type
self.calib_data_loader = calib_data_loader
......@@ -327,7 +327,7 @@ class ModelSpeedupTensorRT(BaseModelSpeedup):
calib = calibrator.Calibrator(calib_data, self.calibration_cache, self.batchsize, self.calibrate_type)
# build inference engine with calibration
engine = build_engine(onnx_path, self.onnx_config, self.extra_layer_bit, self.strict_datatype, calib)
engine = build_engine(onnx_path, self.onnx_config, self.extra_layer_bits, self.strict_datatype, calib)
return engine.create_execution_context()
def _tensorrt_build_withoutcalib(self, onnx_path):
......@@ -344,7 +344,7 @@ class ModelSpeedupTensorRT(BaseModelSpeedup):
tensorrt.IExecutionContext
Context for executing inference using an ICudaEngine
"""
engine = build_engine(onnx_path, self.onnx_config, self.extra_layer_bit, self.strict_datatype)
engine = build_engine(onnx_path, self.onnx_config, self.extra_layer_bits, self.strict_datatype)
return engine.create_execution_context()
def inference(self, test_data):
......
test/ut/sdk/test_compressor_torch.py
View file @ 86335921
......@@ -49,7 +49,8 @@ class CompressorTestCase(TestCase):
}]
model.relu = torch.nn.ReLU()
quantizer = torch_quantizer.QAT_Quantizer(model, config_list)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
quantizer = torch_quantizer.QAT_Quantizer(model, config_list, optimizer)
quantizer.compress()
modules_to_compress = quantizer.get_modules_to_compress()
modules_to_compress_name = [t[0].name for t in modules_to_compress]
......@@ -317,7 +318,9 @@ class CompressorTestCase(TestCase):
'op_types': ['ReLU']
}]
model.relu = torch.nn.ReLU()
quantizer = torch_quantizer.QAT_Quantizer(model, config_list)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
quantizer = torch_quantizer.QAT_Quantizer(model, config_list, optimizer)
quantizer.compress()
# test quantize
......@@ -350,14 +353,14 @@ class CompressorTestCase(TestCase):
eps = 1e-7
x = torch.tensor([[-0.2, 0], [0.1, 0.2]])
out = model.relu(x)
assert math.isclose(model.relu.module.tracked_min_activation, 0, abs_tol=eps)
assert math.isclose(model.relu.module.tracked_max_activation, 0.002, abs_tol=eps)
assert math.isclose(model.relu.module.tracked_min_output, 0, abs_tol=eps)
assert math.isclose(model.relu.module.tracked_max_output, 0.002, abs_tol=eps)
quantizer.step_with_optimizer()
x = torch.tensor([[0.2, 0.4], [0.6, 0.8]])
out = model.relu(x)
assert math.isclose(model.relu.module.tracked_min_activation, 0.002, abs_tol=eps)
assert math.isclose(model.relu.module.tracked_max_activation, 0.00998, abs_tol=eps)
assert math.isclose(model.relu.module.tracked_min_output, 0.002, abs_tol=eps)
assert math.isclose(model.relu.module.tracked_max_output, 0.00998, abs_tol=eps)
def test_torch_quantizer_export(self):
config_list_qat = [{
......@@ -392,7 +395,8 @@ class CompressorTestCase(TestCase):
for config, quantize_algorithm in zip(config_set, quantize_algorithm_set):
model = TorchModel()
model.relu = torch.nn.ReLU()
quantizer = quantize_algorithm(model, config)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
quantizer = quantize_algorithm(model, config, optimizer)
quantizer.compress()
x = torch.rand((1, 1, 28, 28), requires_grad=True)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment