• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# mypy: allow-untyped-defs
2from .fake_quantize import *  # noqa: F403
3from .fuse_modules import fuse_modules
4from .fuser_method_mappings import *  # noqa: F403
5from .observer import *  # noqa: F403
6from .qconfig import *  # noqa: F403
7from .quant_type import *  # noqa: F403
8from .quantization_mappings import *  # noqa: F403
9from .quantize import *  # noqa: F403
10from .quantize_jit import *  # noqa: F403
11from .stubs import *  # noqa: F403
12
13
14def default_eval_fn(model, calib_data):
15    r"""
16    Default evaluation function takes a torch.utils.data.Dataset or a list of
17    input Tensors and run the model on the dataset
18    """
19    for data, target in calib_data:
20        model(data)
21
22
23__all__ = [
24    "QuantWrapper",
25    "QuantStub",
26    "DeQuantStub",
27    # Top level API for eager mode quantization
28    "quantize",
29    "quantize_dynamic",
30    "quantize_qat",
31    "prepare",
32    "convert",
33    "prepare_qat",
34    # Top level API for graph mode quantization on TorchScript
35    "quantize_jit",
36    "quantize_dynamic_jit",
37    "_prepare_ondevice_dynamic_jit",
38    "_convert_ondevice_dynamic_jit",
39    "_quantize_ondevice_dynamic_jit",
40    # Top level API for graph mode quantization on GraphModule(torch.fx)
41    # 'fuse_fx', 'quantize_fx',  # TODO: add quantize_dynamic_fx
42    # 'prepare_fx', 'prepare_dynamic_fx', 'convert_fx',
43    "QuantType",  # quantization type
44    # custom module APIs
45    "get_default_static_quant_module_mappings",
46    "get_static_quant_module_class",
47    "get_default_dynamic_quant_module_mappings",
48    "get_default_qat_module_mappings",
49    "get_default_qconfig_propagation_list",
50    "get_default_compare_output_module_list",
51    "get_quantized_operator",
52    "get_fuser_method",
53    # Sub functions for `prepare` and `swap_module`
54    "propagate_qconfig_",
55    "add_quant_dequant",
56    "swap_module",
57    "default_eval_fn",
58    # Observers
59    "ObserverBase",
60    "WeightObserver",
61    "HistogramObserver",
62    "observer",
63    "default_observer",
64    "default_weight_observer",
65    "default_placeholder_observer",
66    "default_per_channel_weight_observer",
67    # FakeQuantize (for qat)
68    "default_fake_quant",
69    "default_weight_fake_quant",
70    "default_fixed_qparams_range_neg1to1_fake_quant",
71    "default_fixed_qparams_range_0to1_fake_quant",
72    "default_per_channel_weight_fake_quant",
73    "default_histogram_fake_quant",
74    # QConfig
75    "QConfig",
76    "default_qconfig",
77    "default_dynamic_qconfig",
78    "float16_dynamic_qconfig",
79    "float_qparams_weight_only_qconfig",
80    # QAT utilities
81    "default_qat_qconfig",
82    "prepare_qat",
83    "quantize_qat",
84    # module transformations
85    "fuse_modules",
86]
87