-
Notifications
You must be signed in to change notification settings - Fork 21.4k
/
_numeric_suite.py
467 lines (385 loc) · 16.9 KB
/
_numeric_suite.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
import torch
import torch.nn as nn
import torch.nn.quantized as nnq
import torch.nn.quantized.dynamic as nnqd
from torch.quantization import prepare
from typing import Dict
from .quantization_mappings import (
get_default_compare_output_module_list,
)
NON_LEAF_MODULE_TO_ADD_OBSERVER_ALLOW_LIST = {
nnqd.Linear,
nnq.Linear,
nnqd.LSTM,
nn.LSTM,
}
def _find_match(str_list, key_str, postfix):
split_str = key_str.split(".")
if split_str[-1] == postfix:
match_string = "".join(key_str.split(".")[0:-1])
for s2 in str_list:
pattern1 = "".join(s2.split(".")[0:-1])
pattern2 = "".join(s2.split(".")[0:-2])
if match_string == pattern1:
return s2
if match_string == pattern2:
return s2
# For matching "fc.weight" and "fc._packed_params._packed_params"
if postfix == "_packed_params":
match_string = "".join(key_str.split(".")[0:-2])
if len(match_string) == 0:
return None
for s2 in str_list:
pattern1 = "".join(s2.split(".")[0:-1])
pattern2 = "".join(s2.split(".")[0:-2])
if match_string == pattern1:
return s2
if match_string == pattern2:
return s2
else:
return None
def compare_weights(float_dict, quantized_dict):
r"""Compare the weights of the float module with its corresponding quantized
module. Return a dict with key corresponding to module names and each entry being
a dictionary with two keys 'float' and 'quantized', containing the float and
quantized weights. This dict can be used to compare and compute the quantization
error of the weights of float and quantized models.
Example usage:
wt_compare_dict = compare_weights(float_model.state_dict(), qmodel.state_dict())
for key in wt_compare_dict:
print(key, compute_error(wt_compare_dict[key]['float'], wt_compare_dict[key]['quantized'].dequantize()))
Args:
float_dict: state dict of the float model
quantized_dict: state dict of the quantized model
Return:
weight_dict: dict with key corresponding to module names and each entry being
a dictionary with two keys 'float' and 'quantized', containing the float and
quantized weights
"""
torch._C._log_api_usage_once("quantization_api._numeric_suite.compare_weights")
weight_dict: Dict[str, Dict] = {}
for key in quantized_dict:
match_key = _find_match(float_dict, key, "weight")
if match_key is not None:
weight_dict[key] = {}
weight_dict[key]["float"] = float_dict[match_key]
weight_dict[key]["quantized"] = quantized_dict[key]
continue
# For matching "fc.weight" and "fc._packed_params._packed_params"
match_key = _find_match(float_dict, key, "_packed_params")
if match_key is not None:
weight_dict[key] = {}
weight_dict[key]["float"] = float_dict[match_key]
weight_dict[key]["quantized"] = quantized_dict[key][0]
# For LSTM
split_str = key.split(".")
if split_str[-1] == "param" and split_str[-3] == "_all_weight_values":
layer = split_str[-2]
module_name = ".".join(split_str[:-3])
float_weight_ih_key = module_name + ".weight_ih_l" + layer
float_weight_hh_key = module_name + ".weight_hh_l" + layer
if float_weight_ih_key in float_dict and float_weight_hh_key in float_dict:
weight_dict[key] = {}
weight_dict[key]["float"] = float_dict[float_weight_ih_key]
weight_dict[key]["quantized"] = (
quantized_dict[key].__getstate__()[0][4][0].__getstate__()[0][0]
)
weight_dict[key]["float"] = float_dict[float_weight_hh_key]
weight_dict[key]["quantized"] = (
quantized_dict[key].__getstate__()[0][4][1].__getstate__()[0][0]
)
return weight_dict
def _get_logger_dict_helper(mod, target_dict, prefix=""):
r"""This is the helper function for get_logger_dict
Args:
mod: module we want to save all logger stats
prefix: prefix for the current module
target_dict: the dictionary used to save all logger stats
"""
def get_prefix(prefix):
return prefix if prefix == "" else prefix + "."
for name, child in mod.named_children():
if isinstance(child, Logger):
target_dict[get_prefix(prefix) + "stats"] = child.stats
break
for name, child in mod.named_children():
module_prefix = get_prefix(prefix) + name if prefix else name
_get_logger_dict_helper(child, target_dict, module_prefix)
def get_logger_dict(mod, prefix=""):
r"""Traverse the modules and save all logger stats into target dict.
This is mainly used for quantization accuracy debug.
Type of loggers supported:
ShadowLogger: used to log the outputs of the quantized module and its
matching float shadow module,
OutputLogger: used to log the outputs of the modules
Args:
mod: module we want to save all logger stats
prefix: prefix for the current module
Return:
target_dict: the dictionary used to save all logger stats
"""
torch._C._log_api_usage_once("quantization_api._numeric_suite.get_logger_dict")
target_dict: Dict[str, Dict] = {}
_get_logger_dict_helper(mod, target_dict, prefix)
return target_dict
class Logger(nn.Module):
r"""Base class for stats logging
"""
def __init__(self):
super(Logger, self).__init__()
self.stats = {}
def forward(self, x):
pass
class ShadowLogger(Logger):
r"""Class used in Shadow module to record the outputs of the original and
shadow modules.
"""
def __init__(self):
super(ShadowLogger, self).__init__()
self.stats["float"] = None
self.stats["quantized"] = None
def forward(self, x, y):
if len(x) > 1:
x = x[0]
if len(y) > 1:
y = y[0]
if self.stats["quantized"] is None:
self.stats["quantized"] = x.detach()
else:
self.stats["quantized"] = torch.cat((self.stats["quantized"], x.detach()))
if self.stats["float"] is None:
self.stats["float"] = y.detach()
else:
self.stats["float"] = torch.cat((self.stats["float"], y.detach()))
class OutputLogger(Logger):
r"""Class used to log the outputs of the module
"""
def __init__(self):
super(OutputLogger, self).__init__()
self.stats["tensor_val"] = None
def forward(self, x):
if self.stats["tensor_val"] is None:
self.stats["tensor_val"] = x
else:
self.stats["tensor_val"] = torch.cat((self.stats["tensor_val"], x))
return x
def _convert_tuple_to_list(t):
return list(_convert_tuple_to_list(x) for x in t) if type(t) is tuple else t
def _dequantize_tensor_list(t):
return (
list(_dequantize_tensor_list(x) for x in t)
if type(t) is list
else t.dequantize()
if t.is_quantized
else t
)
class Shadow(nn.Module):
r"""Shadow module attaches the float module to its matching quantized module
as the shadow. Then it uses Logger module to process the outputs of both
modules.
Args:
q_module: module quantized from float_module that we want to shadow
float_module: float module used to shadow q_module
Logger: type of logger used to process the outputs of q_module and
float_module. ShadowLogger or custom loggers can be used.
"""
def __init__(self, q_module, float_module, Logger):
super(Shadow, self).__init__()
self.orig_module = q_module
self.shadow_module = float_module
self.dequant = nnq.DeQuantize()
self.logger = Logger()
def forward(self, *x):
xl = _convert_tuple_to_list(x)
output = self.orig_module(*xl)
xl_float = _dequantize_tensor_list(xl)
shadow_output = self.shadow_module(*xl_float)
self.logger(output, shadow_output)
return output
def add(self, x, y):
output = self.orig_module.add(x, y)
x = x.dequantize()
y = y.dequantize()
shadow_output = self.shadow_module.add(x, y)
self.logger(output, shadow_output)
return output
def add_scalar(self, x, y):
output = self.orig_module.add_scalar(x, y)
x = x.dequantize()
shadow_output = self.shadow_module.add_scalar(x, y)
self.logger(output, shadow_output)
return output
def mul(self, x, y):
output = self.orig_module.mul(x, y)
x = x.dequantize()
y = y.dequantize()
shadow_output = self.shadow_module.mul(x, y)
self.logger(output, shadow_output)
return output
def mul_scalar(self, x, y):
output = self.orig_module.mul_scalar(x, y)
x = x.dequantize()
shadow_output = self.shadow_module.mul_scalar(x, y)
self.logger(output, shadow_output)
return output
def cat(self, x, dim=0):
output = self.orig_module.cat(x, dim)
x = [y.dequantize() for y in x]
shadow_output = self.shadow_module.cat(x, dim)
self.logger(output, shadow_output)
return output
def add_relu(self, x, y):
output = self.orig_module.add_relu(x, y)
x = x.dequantize()
y = y.dequantize()
shadow_output = self.shadow_module.add_relu(x, y)
self.logger(output, shadow_output)
return output
def prepare_model_with_stubs(float_module, q_module, module_swap_list, Logger):
r"""Prepare the model by attaching the float module to its matching quantized
module as the shadow if the float module type is in module_swap_list.
Example usage:
prepare_model_with_stubs(float_model, q_model, module_swap_list, Logger)
q_model(data)
ob_dict = get_logger_dict(q_model)
Args:
float_module: float module used to generate the q_module
q_module: module quantized from float_module
module_swap_list: list of float module types to attach the shadow
Logger: type of logger to be used in shadow module to process the outputs of
quantized module and its float shadow module
"""
torch._C._log_api_usage_once("quantization_api._numeric_suite.prepare_model_with_stubs")
float_module_children = {}
for name, mod in float_module.named_children():
float_module_children[name] = mod
reassign = {}
for name, mod in q_module.named_children():
if name not in float_module_children:
continue
float_mod = float_module_children[name]
if type(float_mod) not in module_swap_list:
prepare_model_with_stubs(float_mod, mod, module_swap_list, Logger)
if type(float_mod) in module_swap_list:
reassign[name] = Shadow(mod, float_mod, Logger)
for key, value in reassign.items():
q_module._modules[key] = value
def compare_model_stub(
float_model, q_model, module_swap_list, *data, Logger=ShadowLogger
):
r"""Compare quantized module in a model with its floating point counterpart,
feeding both of them the same input. Return a dict with key corresponding to
module names and each entry being a dictionary with two keys 'float' and
'quantized', containing the output tensors of quantized and its matching
float shadow module. This dict can be used to compare and compute the module
level quantization error.
This function first call prepare_model_with_stubs() to swap the quantized
module that we want to compare with the Shadow module, which takes quantized
module, corresponding float module and logger as input, and creates a forward
path inside to make the float module to shadow quantized module sharing the
same input. The logger can be customizable, default logger is ShadowLogger
and it will save the outputs of the quantized module and float module that
can be used to compute the module level quantization error.
Example usage:
module_swap_list = [torchvision.models.quantization.resnet.QuantizableBasicBlock]
ob_dict = compare_model_stub(float_model,qmodel,module_swap_list, data)
for key in ob_dict:
print(key, compute_error(ob_dict[key]['float'], ob_dict[key]['quantized'].dequantize()))
Args:
float_model: float model used to generate the q_model
q_model: model quantized from float_model
module_swap_list: list of float module types at which shadow modules will
be attached.
data: input data used to run the prepared q_model
Logger: type of logger to be used in shadow module to process the outputs of
quantized module and its float shadow module
"""
torch._C._log_api_usage_once("quantization_api._numeric_suite.compare_model_stub")
prepare_model_with_stubs(float_model, q_model, module_swap_list, Logger)
q_model(*data)
ob_dict = get_logger_dict(q_model)
return ob_dict
def get_matching_activations(float_module, q_module):
r"""Find the matching activation between float and quantized modules.
Args:
float_module: float module used to generate the q_module
q_module: module quantized from float_module
Return:
act_dict: dict with key corresponding to quantized module names and each
entry being a dictionary with two keys 'float' and 'quantized', containing
the matching float and quantized activations
"""
torch._C._log_api_usage_once("quantization_api._numeric_suite.get_matching_activations")
float_dict = get_logger_dict(float_module)
quantized_dict = get_logger_dict(q_module)
act_dict: Dict[str, Dict] = {}
for key in quantized_dict:
match_key = _find_match(sorted(float_dict, reverse=True), key, "stats")
if match_key is not None:
act_dict[key] = {}
act_dict[key]["float"] = float_dict[match_key]["tensor_val"]
act_dict[key]["quantized"] = quantized_dict[key]["tensor_val"]
return act_dict
def prepare_model_outputs(
float_module,
q_module,
Logger=OutputLogger,
allow_list=None
):
r"""Prepare the model by attaching the logger to both float module
and quantized module if they are in the allow_list.
Args:
float_module: float module used to generate the q_module
q_module: module quantized from float_module
Logger: type of logger to be attached to float_module and q_module
allow_list: list of module types to attach logger
"""
torch._C._log_api_usage_once("quantization_api._numeric_suite.prepare_model_outputs")
if allow_list is None:
allow_list = get_default_compare_output_module_list()
qconfig_debug = torch.quantization.QConfig(activation=Logger, weight=None)
float_module.qconfig = qconfig_debug
prepare(float_module, inplace=True, allow_list=allow_list)
q_module.qconfig = qconfig_debug
prepare(
q_module,
inplace=True,
allow_list=allow_list,
observer_non_leaf_module_list=NON_LEAF_MODULE_TO_ADD_OBSERVER_ALLOW_LIST,
)
def compare_model_outputs(
float_model,
q_model,
*data,
Logger=OutputLogger,
allow_list=None
):
r"""Compare output activations between float and quantized models at
corresponding locations for the same input. Return a dict with key corresponding
to quantized module names and each entry being a dictionary with two keys
'float' and 'quantized', containing the activations of quantized model and
float model at matching locations. This dict can be used to compare and
compute the propagation quantization error.
Example usage:
act_compare_dict = compare_model_outputs(float_model, qmodel, data)
for key in act_compare_dict:
print(key, compute_error(act_compare_dict[key]['float'], act_compare_dict[key]['quantized'].dequantize()))
Args:
float_model: float model used to generate the q_model
q_model: model quantized from float_model
data: input data used to run the prepared float_model and q_model
Logger: type of logger to be attached to float_module and q_module
allow_list: list of module types to attach logger
Return:
act_compare_dict: dict with key corresponding to quantized module names
and each entry being a dictionary with two keys 'float' and 'quantized',
containing the matching float and quantized activations
"""
torch._C._log_api_usage_once("quantization_api._numeric_suite.compare_model_outputs")
if allow_list is None:
allow_list = get_default_compare_output_module_list()
prepare_model_outputs(float_model, q_model, Logger, allow_list)
float_model(*data)
q_model(*data)
act_compare_dict = get_matching_activations(float_model, q_model)
return act_compare_dict