From 4292dfbb2b7a71d27a90a7863388bb693987ac2b Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Thu, 7 Oct 2021 09:19:26 +0100 Subject: [PATCH] use torch.ao.quantization --- .../classification/train_quantization.py | 20 +++++++++---------- references/classification/utils.py | 4 ++-- test/test_models.py | 10 +++++----- torchvision/models/quantization/googlenet.py | 6 +++--- torchvision/models/quantization/inception.py | 6 +++--- .../models/quantization/mobilenetv2.py | 2 +- .../models/quantization/mobilenetv3.py | 8 ++++---- torchvision/models/quantization/resnet.py | 12 +++++------ .../models/quantization/shufflenetv2.py | 10 +++++----- torchvision/models/quantization/utils.py | 14 ++++++------- 10 files changed, 46 insertions(+), 46 deletions(-) diff --git a/references/classification/train_quantization.py b/references/classification/train_quantization.py index 5bf64aea721..72598999952 100644 --- a/references/classification/train_quantization.py +++ b/references/classification/train_quantization.py @@ -4,7 +4,7 @@ import time import torch -import torch.quantization +import torch.ao.quantization import torch.utils.data import torchvision import utils @@ -51,8 +51,8 @@ def main(args): if not (args.test_only or args.post_training_quantize): model.fuse_model() - model.qconfig = torch.quantization.get_default_qat_qconfig(args.backend) - torch.quantization.prepare_qat(model, inplace=True) + model.qconfig = torch.ao.quantization.get_default_qat_qconfig(args.backend) + torch.ao.quantization.prepare_qat(model, inplace=True) if args.distributed and args.sync_bn: model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) @@ -85,12 +85,12 @@ def main(args): ) model.eval() model.fuse_model() - model.qconfig = torch.quantization.get_default_qconfig(args.backend) - torch.quantization.prepare(model, inplace=True) + model.qconfig = torch.ao.quantization.get_default_qconfig(args.backend) + torch.ao.quantization.prepare(model, inplace=True) # Calibrate first print("Calibrating") evaluate(model, criterion, data_loader_calibration, device=device, print_freq=1) - torch.quantization.convert(model, inplace=True) + torch.ao.quantization.convert(model, inplace=True) if args.output_dir: print("Saving quantized model") if utils.is_main_process(): @@ -103,8 +103,8 @@ def main(args): evaluate(model, criterion, data_loader_test, device=device) return - model.apply(torch.quantization.enable_observer) - model.apply(torch.quantization.enable_fake_quant) + model.apply(torch.ao.quantization.enable_observer) + model.apply(torch.ao.quantization.enable_fake_quant) start_time = time.time() for epoch in range(args.start_epoch, args.epochs): if args.distributed: @@ -115,7 +115,7 @@ def main(args): with torch.no_grad(): if epoch >= args.num_observer_update_epochs: print("Disabling observer for subseq epochs, epoch = ", epoch) - model.apply(torch.quantization.disable_observer) + model.apply(torch.ao.quantization.disable_observer) if epoch >= args.num_batch_norm_update_epochs: print("Freezing BN for subseq epochs, epoch = ", epoch) model.apply(torch.nn.intrinsic.qat.freeze_bn_stats) @@ -125,7 +125,7 @@ def main(args): quantized_eval_model = copy.deepcopy(model_without_ddp) quantized_eval_model.eval() quantized_eval_model.to(torch.device("cpu")) - torch.quantization.convert(quantized_eval_model, inplace=True) + torch.ao.quantization.convert(quantized_eval_model, inplace=True) print("Evaluate Quantized model") evaluate(quantized_eval_model, criterion, data_loader_test, device=torch.device("cpu")) diff --git a/references/classification/utils.py b/references/classification/utils.py index 5dbb6b8fd24..52506bf6636 100644 --- a/references/classification/utils.py +++ b/references/classification/utils.py @@ -350,8 +350,8 @@ def store_model_weights(model, checkpoint_path, checkpoint_key="model", strict=T # Quantized Classification model = M.quantization.mobilenet_v3_large(pretrained=False, quantize=False) model.fuse_model() - model.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack') - _ = torch.quantization.prepare_qat(model, inplace=True) + model.qconfig = torch.ao.quantization.get_default_qat_qconfig('qnnpack') + _ = torch.ao.quantization.prepare_qat(model, inplace=True) print(store_model_weights(model, './qat.pth')) # Object Detection diff --git a/test/test_models.py b/test/test_models.py index 5e5b3429778..84fd98f5076 100644 --- a/test/test_models.py +++ b/test/test_models.py @@ -705,19 +705,19 @@ def test_quantized_classification_model(model_name): model = torchvision.models.quantization.__dict__[model_name](**kwargs) if eval_mode: model.eval() - model.qconfig = torch.quantization.default_qconfig + model.qconfig = torch.ao.quantization.default_qconfig else: model.train() - model.qconfig = torch.quantization.default_qat_qconfig + model.qconfig = torch.ao.quantization.default_qat_qconfig model.fuse_model() if eval_mode: - torch.quantization.prepare(model, inplace=True) + torch.ao.quantization.prepare(model, inplace=True) else: - torch.quantization.prepare_qat(model, inplace=True) + torch.ao.quantization.prepare_qat(model, inplace=True) model.eval() - torch.quantization.convert(model, inplace=True) + torch.ao.quantization.convert(model, inplace=True) try: torch.jit.script(model) diff --git a/torchvision/models/quantization/googlenet.py b/torchvision/models/quantization/googlenet.py index 4b6d25e013c..80261aaf295 100644 --- a/torchvision/models/quantization/googlenet.py +++ b/torchvision/models/quantization/googlenet.py @@ -94,7 +94,7 @@ def forward(self, x: Tensor) -> Tensor: return x def fuse_model(self) -> None: - torch.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True) + torch.ao.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True) class QuantizableInception(Inception): @@ -142,8 +142,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super(QuantizableGoogLeNet, self).__init__( # type: ignore[misc] blocks=[QuantizableBasicConv2d, QuantizableInception, QuantizableInceptionAux], *args, **kwargs ) - self.quant = torch.quantization.QuantStub() - self.dequant = torch.quantization.DeQuantStub() + self.quant = torch.ao.quantization.QuantStub() + self.dequant = torch.ao.quantization.DeQuantStub() def forward(self, x: Tensor) -> GoogLeNetOutputs: x = self._transform_input(x) diff --git a/torchvision/models/quantization/inception.py b/torchvision/models/quantization/inception.py index acad3f6df53..f1590d50b84 100644 --- a/torchvision/models/quantization/inception.py +++ b/torchvision/models/quantization/inception.py @@ -102,7 +102,7 @@ def forward(self, x: Tensor) -> Tensor: return x def fuse_model(self) -> None: - torch.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True) + torch.ao.quantization.fuse_modules(self, ["conv", "bn", "relu"], inplace=True) class QuantizableInceptionA(inception_module.InceptionA): @@ -222,8 +222,8 @@ def __init__( QuantizableInceptionAux, ], ) - self.quant = torch.quantization.QuantStub() - self.dequant = torch.quantization.DeQuantStub() + self.quant = torch.ao.quantization.QuantStub() + self.dequant = torch.ao.quantization.DeQuantStub() def forward(self, x: Tensor) -> InceptionOutputs: x = self._transform_input(x) diff --git a/torchvision/models/quantization/mobilenetv2.py b/torchvision/models/quantization/mobilenetv2.py index a2c88cdd388..6da60e9cd02 100644 --- a/torchvision/models/quantization/mobilenetv2.py +++ b/torchvision/models/quantization/mobilenetv2.py @@ -2,7 +2,7 @@ from torch import Tensor from torch import nn -from torch.quantization import QuantStub, DeQuantStub, fuse_modules +from torch.ao.quantization import QuantStub, DeQuantStub, fuse_modules from torchvision.models.mobilenetv2 import InvertedResidual, MobileNetV2, model_urls from ..._internally_replaced_utils import load_state_dict_from_url diff --git a/torchvision/models/quantization/mobilenetv3.py b/torchvision/models/quantization/mobilenetv3.py index ad195d178c7..1b98c76fc18 100644 --- a/torchvision/models/quantization/mobilenetv3.py +++ b/torchvision/models/quantization/mobilenetv3.py @@ -2,7 +2,7 @@ import torch from torch import nn, Tensor -from torch.quantization import QuantStub, DeQuantStub, fuse_modules +from torch.ao.quantization import QuantStub, DeQuantStub, fuse_modules from ..._internally_replaced_utils import load_state_dict_from_url from ...ops.misc import ConvNormActivation, SqueezeExcitation @@ -132,13 +132,13 @@ def _mobilenet_v3_model( backend = "qnnpack" model.fuse_model() - model.qconfig = torch.quantization.get_default_qat_qconfig(backend) - torch.quantization.prepare_qat(model, inplace=True) + model.qconfig = torch.ao.quantization.get_default_qat_qconfig(backend) + torch.ao.quantization.prepare_qat(model, inplace=True) if pretrained: _load_weights(arch, model, quant_model_urls.get(arch + "_" + backend, None), progress) - torch.quantization.convert(model, inplace=True) + torch.ao.quantization.convert(model, inplace=True) model.eval() else: if pretrained: diff --git a/torchvision/models/quantization/resnet.py b/torchvision/models/quantization/resnet.py index f7124798254..9fa117b0863 100644 --- a/torchvision/models/quantization/resnet.py +++ b/torchvision/models/quantization/resnet.py @@ -3,7 +3,7 @@ import torch import torch.nn as nn from torch import Tensor -from torch.quantization import fuse_modules +from torch.ao.quantization import fuse_modules from torchvision.models.resnet import Bottleneck, BasicBlock, ResNet, model_urls from ..._internally_replaced_utils import load_state_dict_from_url @@ -42,9 +42,9 @@ def forward(self, x: Tensor) -> Tensor: return out def fuse_model(self) -> None: - torch.quantization.fuse_modules(self, [["conv1", "bn1", "relu"], ["conv2", "bn2"]], inplace=True) + torch.ao.quantization.fuse_modules(self, [["conv1", "bn1", "relu"], ["conv2", "bn2"]], inplace=True) if self.downsample: - torch.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True) + torch.ao.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True) class QuantizableBottleneck(Bottleneck): @@ -75,15 +75,15 @@ def forward(self, x: Tensor) -> Tensor: def fuse_model(self) -> None: fuse_modules(self, [["conv1", "bn1", "relu1"], ["conv2", "bn2", "relu2"], ["conv3", "bn3"]], inplace=True) if self.downsample: - torch.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True) + torch.ao.quantization.fuse_modules(self.downsample, ["0", "1"], inplace=True) class QuantizableResNet(ResNet): def __init__(self, *args: Any, **kwargs: Any) -> None: super(QuantizableResNet, self).__init__(*args, **kwargs) - self.quant = torch.quantization.QuantStub() - self.dequant = torch.quantization.DeQuantStub() + self.quant = torch.ao.quantization.QuantStub() + self.dequant = torch.ao.quantization.DeQuantStub() def forward(self, x: Tensor) -> Tensor: x = self.quant(x) diff --git a/torchvision/models/quantization/shufflenetv2.py b/torchvision/models/quantization/shufflenetv2.py index a4c4aede665..9f5e13d69f4 100644 --- a/torchvision/models/quantization/shufflenetv2.py +++ b/torchvision/models/quantization/shufflenetv2.py @@ -47,8 +47,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super(QuantizableShuffleNetV2, self).__init__( # type: ignore[misc] *args, inverted_residual=QuantizableInvertedResidual, **kwargs ) - self.quant = torch.quantization.QuantStub() - self.dequant = torch.quantization.DeQuantStub() + self.quant = torch.ao.quantization.QuantStub() + self.dequant = torch.ao.quantization.DeQuantStub() def forward(self, x: Tensor) -> Tensor: x = self.quant(x) @@ -66,12 +66,12 @@ def fuse_model(self) -> None: for name, m in self._modules.items(): if name in ["conv1", "conv5"]: - torch.quantization.fuse_modules(m, [["0", "1", "2"]], inplace=True) + torch.ao.quantization.fuse_modules(m, [["0", "1", "2"]], inplace=True) for m in self.modules(): if type(m) == QuantizableInvertedResidual: if len(m.branch1._modules.items()) > 0: - torch.quantization.fuse_modules(m.branch1, [["0", "1"], ["2", "3", "4"]], inplace=True) - torch.quantization.fuse_modules( + torch.ao.quantization.fuse_modules(m.branch1, [["0", "1"], ["2", "3", "4"]], inplace=True) + torch.ao.quantization.fuse_modules( m.branch2, [["0", "1", "2"], ["3", "4"], ["5", "6", "7"]], inplace=True, diff --git a/torchvision/models/quantization/utils.py b/torchvision/models/quantization/utils.py index 74a8287030b..b1fcd6fee5d 100644 --- a/torchvision/models/quantization/utils.py +++ b/torchvision/models/quantization/utils.py @@ -24,19 +24,19 @@ def quantize_model(model: nn.Module, backend: str) -> None: model.eval() # Make sure that weight qconfig matches that of the serialized models if backend == "fbgemm": - model.qconfig = torch.quantization.QConfig( # type: ignore[assignment] - activation=torch.quantization.default_observer, - weight=torch.quantization.default_per_channel_weight_observer, + model.qconfig = torch.ao.quantization.QConfig( # type: ignore[assignment] + activation=torch.ao.quantization.default_observer, + weight=torch.ao.quantization.default_per_channel_weight_observer, ) elif backend == "qnnpack": - model.qconfig = torch.quantization.QConfig( # type: ignore[assignment] - activation=torch.quantization.default_observer, weight=torch.quantization.default_weight_observer + model.qconfig = torch.ao.quantization.QConfig( # type: ignore[assignment] + activation=torch.ao.quantization.default_observer, weight=torch.ao.quantization.default_weight_observer ) # TODO https://github.com/pytorch/vision/pull/4232#pullrequestreview-730461659 model.fuse_model() # type: ignore[operator] - torch.quantization.prepare(model, inplace=True) + torch.ao.quantization.prepare(model, inplace=True) model(_dummy_input_data) - torch.quantization.convert(model, inplace=True) + torch.ao.quantization.convert(model, inplace=True) return