Skip to content

Commit 12386c9

Browse files
committed
fix export code/text
1 parent 018a7dc commit 12386c9

File tree

1 file changed

+10
-9
lines changed

1 file changed

+10
-9
lines changed

prototype_source/pt2e_quant_ptq_static.rst

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -22,27 +22,27 @@ this:
2222
\ /
2323
\ /
2424
—-------------------------------------------------------
25-
| Dynamo Export |
25+
| Export |
2626
—-------------------------------------------------------
2727
|
2828
FX Graph in ATen XNNPACKQuantizer,
2929
| or X86InductorQuantizer,
3030
| or <Other Backend Quantizer>
3131
| /
3232
—--------------------------------------------------------
33-
| prepare_pt2e |
33+
| prepare_pt2e |
3434
—--------------------------------------------------------
3535
|
3636
Calibrate/Train
3737
|
3838
—--------------------------------------------------------
39-
| convert_pt2e |
39+
| convert_pt2e |
4040
—--------------------------------------------------------
4141
|
4242
Reference Quantized Model
4343
|
4444
—--------------------------------------------------------
45-
| Lowering |
45+
| Lowering |
4646
—--------------------------------------------------------
4747
|
4848
Executorch, or Inductor, or <Other Backends>
@@ -189,8 +189,6 @@ and rename it to ``data/resnet18_pretrained_float.pth``.
189189
import numpy as np
190190
191191
import torch
192-
from torch.ao.quantization import get_default_qconfig, QConfigMapping
193-
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx, fuse_fx
194192
import torch.nn as nn
195193
from torch.utils.data import DataLoader
196194
@@ -358,7 +356,10 @@ Here is how you can use ``torch.export`` to export the model:
358356
from torch._export import capture_pre_autograd_graph
359357
360358
example_inputs = (torch.rand(2, 3, 224, 224),)
361-
exported_model, _ = capture_pre_autograd_graph(model_to_quantize, *example_inputs)
359+
exported_model = capture_pre_autograd_graph(model_to_quantize, example_inputs)
360+
# or capture with dynamic dimensions
361+
# from torch._export import dynamic_dim
362+
# exported_model = capture_pre_autograd_graph(model_to_quantize, example_inputs, constraints=[dynamic_dim(example_inputs[0], 0)])
362363
363364
364365
``capture_pre_autograd_graph`` is a short term API, it will be updated to use the offical ``torch.export`` API when that is ready.
@@ -532,9 +533,9 @@ We'll show how to save and load the quantized model.
532533
# Rerun all steps to get a quantized model
533534
model_to_quantize = load_model(saved_model_dir + float_model_file).to("cpu")
534535
model_to_quantize.eval()
535-
import torch._dynamo as torchdynamo
536+
from torch._export import capture_pre_autograd_graph
536537
537-
exported_model, _ = torchdynamo.export(model_to_quantize, *copy.deepcopy(example_inputs), aten_graph=True, tracing_mode="symbolic")
538+
exported_model = capture_pre_autograd_graph(model_to_quantize, example_inputs)
538539
from torch.ao.quantization.quantizer.xnnpack_quantizer import (
539540
XNNPACKQuantizer,
540541
get_symmetric_quantization_config,

0 commit comments

Comments
 (0)