Skip to content

Commit e234e00

Browse files
committed
docs: fixed faulty code snippets in nxp-quantization
1 parent 0129dc9 commit e234e00

1 file changed

Lines changed: 25 additions & 7 deletions

File tree

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

Lines changed: 25 additions & 7 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

@@ -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)