diff --git a/docs/source/backends/nxp/nxp-quantization.md b/docs/source/backends/nxp/nxp-quantization.md index 095f9e29e52..4eb9ac02b2c 100644 --- a/docs/source/backends/nxp/nxp-quantization.md +++ b/docs/source/backends/nxp/nxp-quantization.md @@ -54,11 +54,18 @@ To quantize the model, you can use the PT2E workflow: import torch import torchvision.models as models from torchvision.models.mobilenetv2 import MobileNet_V2_Weights + from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec from executorch.exir import to_edge_transform_and_lower + +# Imported for side effects: registers the quantized out-variant kernels +# so `to_executorch()` can find them. +import executorch.extension.pybindings.portable_lib # noqa: F401 +import executorch.kernels.quantized # noqa: F401 + from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e model = models.mobilenetv2.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).eval() @@ -82,7 +89,10 @@ compile_spec = generate_neutron_compile_spec( et_program = to_edge_transform_and_lower( # (6) torch.export.export(quantized_model, sample_inputs), - partitioner=[NeutronPartitioner(compile_spec=compile_spec)], + partitioner=[NeutronPartitioner( + compile_spec=compile_spec, + neutron_target_spec=neutron_target_spec + )], ).to_executorch() ``` @@ -102,7 +112,7 @@ quantized_graph_module = calibrate_and_quantize( ) ``` -See [PyTorch 2 Export Post Training Quantization](https://docs.pytorch.org/ao/main/tutorials_source/pt2e_quant_ptq.html) for more information. +See [PyTorch 2 Export Post Training Quantization](https://docs.pytorch.org/ao/stable/pt2e_quantization/pt2e_quant_ptq.html) for more information. ### Quantization Aware Training @@ -138,6 +148,7 @@ import torch from torch.utils.data import DataLoader import torchvision.models as models import torchvision.datasets as datasets +import torchvision.transforms as transforms from torchvision.models.mobilenetv2 import MobileNet_V2_Weights from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec @@ -164,10 +175,22 @@ prepared_model = move_exported_model_to_train(prepared_model) # (4) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(prepared_model.parameters(), lr=1e-2, momentum=0.9) -train_data = datasets.ImageNet("./", split="train", transform=...) +transform = transforms.Compose( + [ + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ), + ] +) + +train_data = datasets.ImageNet("./", split="train", transform=transform) train_loader = DataLoader(train_data, batch_size=5) # Training replaces calibration in QAT +num_epochs = 5 for epoch in range(num_epochs): for imgs, labels in train_loader: optimizer.zero_grad() @@ -185,11 +208,6 @@ for epoch in range(num_epochs): prepared_model = move_exported_model_to_eval(prepared_model) # (6) quantized_model = convert_pt2e(prepared_model) # (7) -# Optional step - fixes biasless convolution (see Known Limitations of QAT) -quantized_model = QuantizeFusedConvBnBiasAtenPass( - default_zero_bias=True -)(quantized_model).graph_module - ... ```