@@ -54,11 +54,18 @@ To quantize the model, you can use the PT2E workflow:
5454import torch
5555import torchvision.models as models
5656from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
57+
5758from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
5859from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
5960from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
6061from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec
6162from executorch.exir import to_edge_transform_and_lower
63+
64+ # Imported for side effects: registers the quantized out-variant kernels
65+ # so `to_executorch()` can find them.
66+ import executorch.extension.pybindings.portable_lib # noqa: F401
67+ import executorch.kernels.quantized # noqa: F401
68+
6269from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
6370
6471model = models.mobilenetv2.mobilenet_v2(weights = MobileNet_V2_Weights.DEFAULT ).eval()
@@ -82,7 +89,10 @@ compile_spec = generate_neutron_compile_spec(
8289
8390et_program = to_edge_transform_and_lower( # (6)
8491 torch.export.export(quantized_model, sample_inputs),
85- partitioner = [NeutronPartitioner(compile_spec = compile_spec)],
92+ partitioner = [NeutronPartitioner(
93+ compile_spec = compile_spec,
94+ neutron_target_spec = neutron_target_spec
95+ )],
8696).to_executorch()
8797```
8898
@@ -138,6 +148,7 @@ import torch
138148from torch.utils.data import DataLoader
139149import torchvision.models as models
140150import torchvision.datasets as datasets
151+ import torchvision.transforms as transforms
141152from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
142153from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
143154from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
@@ -164,10 +175,22 @@ prepared_model = move_exported_model_to_train(prepared_model) # (4)
164175criterion = torch.nn.CrossEntropyLoss()
165176optimizer = torch.optim.SGD(prepared_model.parameters(), lr = 1e-2 , momentum = 0.9 )
166177
167- train_data = datasets.ImageNet(" ./" , split = " train" , transform = ... )
178+ transform = transforms.Compose(
179+ [
180+ transforms.Resize(256 ),
181+ transforms.CenterCrop(224 ),
182+ transforms.ToTensor(),
183+ transforms.Normalize(
184+ mean = [0.485 , 0.456 , 0.406 ], std = [0.229 , 0.224 , 0.225 ]
185+ ),
186+ ]
187+ )
188+
189+ train_data = datasets.ImageNet(" ./" , split = " train" , transform = transform)
168190train_loader = DataLoader(train_data, batch_size = 5 )
169191
170192# Training replaces calibration in QAT
193+ num_epochs = 5
171194for epoch in range (num_epochs):
172195 for imgs, labels in train_loader:
173196 optimizer.zero_grad()
@@ -185,11 +208,6 @@ for epoch in range(num_epochs):
185208prepared_model = move_exported_model_to_eval(prepared_model) # (6)
186209quantized_model = convert_pt2e(prepared_model) # (7)
187210
188- # Optional step - fixes biasless convolution (see Known Limitations of QAT)
189- quantized_model = QuantizeFusedConvBnBiasAtenPass(
190- default_zero_bias = True
191- )(quantized_model).graph_module
192-
193211...
194212```
195213
0 commit comments