39
39
import copy
40
40
import torch
41
41
import torchvision
42
- import torch .quantization
42
+ import torch .ao . quantization
43
43
import torch .ao .ns ._numeric_suite_fx as ns
44
- import torch .quantization .quantize_fx as quantize_fx
44
+ import torch .ao .quantization .quantize_fx as quantize_fx
45
+ from torch .ao .quantization .qconfig_mapping import get_default_qconfig_mapping
45
46
46
47
import matplotlib .pyplot as plt
47
48
from tabulate import tabulate
@@ -68,25 +69,23 @@ def plot(xdata, ydata, xlabel, ylabel, title):
68
69
mobilenetv2_float = torchvision .models .quantization .mobilenet_v2 (
69
70
pretrained = True , quantize = False ).eval ()
70
71
72
+ # adjust the default qconfig to make the results more interesting to explore
73
+ # 1. turn off quantization for the first couple of layers
74
+ # 2. use MinMaxObserver for `features.17`, this should lead to worse
75
+ # weight SQNR
76
+ qconfig_mapping = get_default_qconfig_mapping ('fbgemm' )\
77
+ .set_module_name ('features.0' , None )\
78
+ .set_module_name ('features.1' , None )\
79
+ .set_module_name ('features.0' , torch .ao .quantization .default_qconfig )
80
+
71
81
# create quantized model
72
- qconfig_dict = {
73
- '' : torch .quantization .get_default_qconfig ('fbgemm' ),
74
- # adjust the qconfig to make the results more interesting to explore
75
- 'module_name' : [
76
- # turn off quantization for the first couple of layers
77
- ('features.0' , None ),
78
- ('features.1' , None ),
79
- # use MinMaxObserver for `features.17`, this should lead to worse
80
- # weight SQNR
81
- ('features.17' , torch .quantization .default_qconfig ),
82
- ]
83
- }
82
+
84
83
# Note: quantization APIs are inplace, so we save a copy of the float model for
85
84
# later comparison to the quantized model. This is done throughout the
86
85
# tutorial.
87
- mobilenetv2_prepared = quantize_fx .prepare_fx (
88
- copy .deepcopy (mobilenetv2_float ), qconfig_dict )
89
86
datum = torch .randn (1 , 3 , 224 , 224 )
87
+ mobilenetv2_prepared = quantize_fx .prepare_fx (
88
+ copy .deepcopy (mobilenetv2_float ), qconfig_mapping , (datum ,))
90
89
mobilenetv2_prepared (datum )
91
90
# Note: there is a long standing issue that we cannot copy.deepcopy a
92
91
# quantized model. Since quantization APIs are inplace and we need to use
0 commit comments