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.