How to prune part of model with QAT pruner

[2730gf]

Describe the issue:
Hello, I have a model that needs to be quantized. This model consists of backbone, neck and head. I only want to quantify the neck part. There are many model parameters here, I don't want to specify each tensor name. Here is the method I use:
quantizer = QAT_Quantizer(model.neck, config_list, optimizer)
But the optimizer saves the parameters of the entire model. Will it have any effect if I use it like this? At the same time, what is the role of optimizer in QAT, and why does QAT need to pass in an optimizer?
If there is something wrong with my approach, what should I do?

Environment:

• NNI version:
• Training service (local|remote|pai|aml|etc):
• Client OS:
• Server OS (for remote mode only):
• Python version:
• PyTorch/TensorFlow version:
• Is conda/virtualenv/venv used?:
• Is running in Docker?:

Configuration:

• Experiment config (remember to remove secrets!):
• Search space:

Log message:

• nnimanager.log:
• dispatcher.log:
• nnictl stdout and stderr:

How to reproduce it?:

[J-shang]

I think it's OK.

optimizer.step will be replaced in QATQuantizer by another step:

def new_step(optimizer):
    optimizer.step()
    model.steps.add_(1)

So QATQuantizer don't need the parameters bounded in your optimizer, just use it to count steps.

But please pay attention to the config_list and the calibration_config, the module name in these dict will be xxx instead of neck.xxx, so if you want to use the calibration_config to do the real speedup (i.e., build a TensorRT engine), remember to add neck. prefix for the module name.

[2730gf]

Thank you for your reply. I am using the mmdetection framework, and the optimizer used is not torch.nn.optim, so an error will be reported. How can I solve this?

[2730gf]

Hello, can anyone provide a solution?

[J-shang]

Thank you for your reply. I am using the mmdetection framework, and the optimizer used is not torch.nn.optim, so an error will be reported. How can I solve this?

could you show the error and the optimizer you used?

[2730gf]

Hello, I am reading the source code to try to solve the problem of optimizer; but I found another problem, if QAT does not specify dummy_input, there will be this bug below, is this normal? I see dummy_input=None can be specified
AssertionError: Could not found shapes for layer conv1

[2730gf]

Add my code here:
quantizer = QAT_Quantizer(model, config_list, optimizer) quantizer.compress()

[J-shang]

Hello, I am reading the source code to try to solve the problem of optimizer; but I found another problem, if QAT does not specify dummy_input, there will be this bug below, is this normal? I see dummy_input=None can be specified AssertionError: Could not found shapes for layer conv1

I check the QAT logic, I think dummy_input should not be None, it's a required parameter for QAT. You could random generate a dummy input by torch.rand(...).to(device)

[2730gf]

The link I refer to is as follows:

nni/examples/model_compress/end2end_compression.py

Line 220 in 75e5d5b

quantizer = QAT_Quantizer(model, config_list, optimizer)

Dummy_input is not used here, and whether the shape of this dummy_input needs to be the same as the input size of the model. If the shape is inconsistent, it will cause the QAT to fail, right?

[J-shang]

The link I refer to is as follows:

nni/examples/model_compress/end2end_compression.py

Line 220 in 75e5d5b

quantizer = QAT_Quantizer(model, config_list, optimizer)

Dummy_input is not used here, and whether the shape of this dummy_input needs to be the same as the input size of the model. If the shape is inconsistent, it will cause the QAT to fail, right?

Seems an old version example, I will update it.

yes, QAT need to know the input/output shape of each quantized layer.