Skip to content

Commit b769ece

Browse files
novak-vaclavpytorchbot
authored andcommitted
NXP backend: fixed docs in nxp-quantization.md (#21464)
### Summary Fixed code snippets in `nxp-quantization.md`. - there was an issue with faulty imports, where some kernels were not registered properly. - added missing `neutron_target_spec` argument. - here was another issue with the snippet for QAT. I fixed the code so it can be copy-pasted and run directly. I also removed the optional biasless convolution pass step. It makes the whole snippet quite misleading, since the example MobileNet does not have biasless convolutions, thus the code fails. I think the whole caveat is understandable from the last section about known limitations. - fixed invalid link. ### Test plan No need to run tests, these changes affect only docs. (cherry picked from commit 85a204a)
1 parent cd380e7 commit b769ece

1 file changed

Lines changed: 26 additions & 8 deletions

File tree

docs/source/backends/nxp/nxp-quantization.md

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -54,11 +54,18 @@ To quantize the model, you can use the PT2E workflow:
5454
import torch
5555
import torchvision.models as models
5656
from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
57+
5758
from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
5859
from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
5960
from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
6061
from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec
6162
from 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+
6269
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
6370

6471
model = models.mobilenetv2.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).eval()
@@ -82,7 +89,10 @@ compile_spec = generate_neutron_compile_spec(
8289

8390
et_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

@@ -102,7 +112,7 @@ quantized_graph_module = calibrate_and_quantize(
102112
)
103113
```
104114

105-
See [PyTorch 2 Export Post Training Quantization](https://docs.pytorch.org/ao/main/tutorials_source/pt2e_quant_ptq.html) for more information.
115+
See [PyTorch 2 Export Post Training Quantization](https://docs.pytorch.org/ao/stable/pt2e_quantization/pt2e_quant_ptq.html) for more information.
106116

107117
### Quantization Aware Training
108118

@@ -138,6 +148,7 @@ import torch
138148
from torch.utils.data import DataLoader
139149
import torchvision.models as models
140150
import torchvision.datasets as datasets
151+
import torchvision.transforms as transforms
141152
from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
142153
from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
143154
from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
@@ -164,10 +175,22 @@ prepared_model = move_exported_model_to_train(prepared_model) # (4)
164175
criterion = torch.nn.CrossEntropyLoss()
165176
optimizer = 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)
168190
train_loader = DataLoader(train_data, batch_size=5)
169191

170192
# Training replaces calibration in QAT
193+
num_epochs = 5
171194
for 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):
185208
prepared_model = move_exported_model_to_eval(prepared_model) # (6)
186209
quantized_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

Comments
 (0)