Saving modified fx graph

HM_COS: <p>I quantized resnet18 with pt2e. In the qunatized model I replaced few patterns with custom functions using subgraph_rewriter.replace_pattern and torch.fx.wrap for custom functions.<br>
How to save the modified graph with custom functions ?<br>
If I use torch.export and save the model, custom functions are replaced with standard operators. I won’t be able to work on it further at later point in time.</p>
<h2><a name="p-461089-code-snippet-is-below-torchfxlib-has-the-custom-functions-which-is-mainly-used-to-fuse-the-qdq-bw-nodes-into-the-same-node-1" class="anchor" href="#p-461089-code-snippet-is-below-torchfxlib-has-the-custom-functions-which-is-mainly-used-to-fuse-the-qdq-bw-nodes-into-the-same-node-1"></a>Code snippet is below. torchfxlib has the custom functions which is mainly used to fuse the q/dq b/w nodes into the same node.</h2>
<p>model_to_quantize_wt = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1).to(“cpu”)<br>
model_to_quantize = resnet.resnet18()<br>
model_to_quantize.load_state_dict(model_to_quantize_wt.state_dict())<br>
model_to_quantize.eval()</p>
<p>exported_model = torch.export.export_for_training(model_to_quantize, example_inputs).module()</p>
<p>from torch.ao.quantization.quantizer.xnnpack_quantizer import (<br>
XNNPACKQuantizer,<br>
get_symmetric_quantization_config,<br>
)</p>
<p>qconfig = get_quantization_config()</p>
<h1><a name="p-461089-select-a-quantizer-and-configure-2" class="anchor" href="#p-461089-select-a-quantizer-and-configure-2"></a>select a quantizer and configure</h1>
<p>quantizer = XNNPACKQuantizer()<br>
quantizer.set_global(qconfig)</p>
<h1><a name="p-461089-prepare-for-quantization-3" class="anchor" href="#p-461089-prepare-for-quantization-3"></a>Prepare for quantization</h1>
<p>prepared_model = prepare_pt2e(exported_model, quantizer)</p>
<p>def calibrate(model, data_loader):<br>
torch.ao.quantization.move_exported_model_to_eval(model)<br>
with torch.no_grad():<br>
for image, target in data_loader:<br>
model(image)</p>
<h1><a name="p-461089-run-calibration-on-sample-data-4" class="anchor" href="#p-461089-run-calibration-on-sample-data-4"></a>run calibration on sample data</h1>
<p>calibrate(prepared_model, data_loader_cal)<br>
quantized_model = convert_pt2e(prepared_model, fold_quantize=True)</p>
<h1><a name="p-461089-transform-the-graph-with-custom-functions-5" class="anchor" href="#p-461089-transform-the-graph-with-custom-functions-5"></a>Transform the graph with custom functions</h1>
<p>matched_pattern = subgraph_rewriter.replace_pattern(quantized_model, torchfxlib.pattern_qConv2dRelu, torchfxlib.replacement_qConv2dRelu)<br>
matched_pattern = subgraph_rewriter.replace_pattern(quantized_model, torchfxlib.pattern_qConv2d, torchfxlib.replacement_qConv2d)<br>
matched_pattern = subgraph_rewriter.replace_pattern(quantized_model, torchfxlib.pattern_qmax_pool2d, torchfxlib.replacement_qmax_pool2d)<br>
matched_pattern = subgraph_rewriter.replace_pattern(quantized_model, torchfxlib.pattern_qadd, torchfxlib.replacement_qadd)<br>
matched_pattern = subgraph_rewriter.replace_pattern(quantized_model, torchfxlib.pattern_qavg_pool2d, torchfxlib.replacement_qavg_pool2d)<br>
matched_pattern = subgraph_rewriter.replace_pattern(quantized_model, torchfxlib.pattern_qflatten, torchfxlib.replacement_qflatten)<br>
matched_pattern = subgraph_rewriter.replace_pattern(quantized_model, torchfxlib.pattern_qlinear, torchfxlib.replacement_qlinear)<br>
matched_pattern = subgraph_rewriter.replace_pattern(quantized_model, torchfxlib.pattern_q, torchfxlib.replacement_q)<br>
matched_pattern = subgraph_rewriter.replace_pattern(quantized_model, torchfxlib.pattern_dq, torchfxlib.replacement_dq)</p>
<h1><a name="p-461089-export-the-transformed-graph-and-save-exportedprogram-6" class="anchor" href="#p-461089-export-the-transformed-graph-and-save-exportedprogram-6"></a>Export the transformed graph and Save ExportedProgram</h1>
<p>pt2e_quantized_model_file_path = saved_model_dir + “resnet18_quantized.pth”<br>
quantized_ep = torch.export.export(quantized_model, example_inputs)<br>
torch.export.save(quantized_ep, pt2e_quantized_model_file_path)</p>