Develop QAT with MQBench
Given a model which has been prepared by mqbench.prepare_by_platform
, fake quantize nodes are accessible for further training. Then, we divide parameters into two groups: normal parameters and quantization parameters:
1# SomeFakeQuantize
2normal_params = []
3quantization_params = []
4for n, m in model.named_modules():
5 if isinstance(m, SomeFakeQuantize):
6 quantization_params.append(m.parameters())
7 else:
8 normal_params.append(m.parameters())
Then get them into optimizer, e.g.:
1# normal_lr
2# quant_lr
3# default_lr
4opt = optim.SGD(
5 [
6 {'params': quantization_params, 'lr': quant_lr},
7 {'params': normal_params, 'lr': normal_lr},
8 ], lr=default_lr
9)
So, as shown above, a QAT model gets a fake quantization module and its related quantization parameters for subsequent quantization aware training. To learn more about how to add a fake quantize, see also Quantization Function.