Add CoreMLQuantizer in coremltools.optimize.torch to support PyTorch Export based quantization

[pulkital]

Adds a CoreMLQuantizer which extends PyTorch 2.0 Export Quantizer to enable CoreML quantization rules.

Usage:

import torch.nn as nn
from torch._export import capture_pre_autograd_graph
from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_qat_pt2e

from coremltools.optimize.torch.quantization._coreml_quantizer import CoreMLQuantizer

model = nn.Sequential(
    OrderedDict(
        {
            "conv": nn.Conv2d(1, 20, (3, 3)),
            "relu1": nn.ReLU(),
            "conv2": nn.Conv2d(20, 20, (3, 3)),
            "relu2": nn.ReLU(),
        }
    )
)

loss_fn = define_loss()

# initialize the annotator with quantization config
config = LinearQuantizerConfig.from_dict(
    {
        "global_config": {
            "quantization_scheme": "symmetric",
        }
    }
)
quantizer = CoreMLQuantizer(config)

example_inputs = torch.randn(1, 1, 28, 28)

# create export graph
exported_model = capture_pre_autograd_graph(model, (example_inputs,))

# prepare the model to insert FakeQuantize layers for QAT
prepared_model = prepare_qat_pt2e(exported_model, quantizer)

# use prepared model in your PyTorch training loop
for inputs, labels in data:
    output = prepared_model(inputs)
    loss = loss_fn(output, labels)
    loss.backward()
    optimizer.step()
    # turn observers/quantizers on/off depending on iteration number

# convert operations to their quanitzed counterparts using parameters learnt via QAT
model = convert_pt2e(prepared_model)

[YifanShenSZ]

Added 2 changes:

1. Fix linter issue
2. Quantization conversion test

Testing: ✅

1. GitLab CI
2. Locally verified backward compatibility with pytorch 1.13.1
3. Locally verified executorch tests

[YifanShenSZ]

Awesome! We are now one step closer to ExecuTorch CoreMLQuantizer! Thanks for the great work!

[jerryzh168]

@pulkital @YifanShenSZ we are deprecating pt2e quantization from pytorch/pytorch (pytorch/pytorch#169151),

wondering if you can help updating the callsites to import from torchao instead?

coremltools/coremltools/optimize/torch/quantization/_coreml_quantizer.py

Lines 12 to 13 in 0434e85

	from torch.ao.quantization.quantizer.quantizer import Quantizer as _TorchQuantizer
	from torch.ao.quantization.quantizer.xnnpack_quantizer import _get_module_name_filter

before

from torch.ao.quantization.quantizer.quantizer import Quantizer as _TorchQuantizer
from torch.ao.quantization.quantizer.xnnpack_quantizer import _get_module_name_filter

after:

from torchao.quantization.pt2e.quantizer import Quantizer as _TorchQuantizer
from torchao.quantization.pt2e.quantizer import get_module_name_filter

[jerryzh168]

please take a look #2634