Update quantization imports for `torch.ao.quantization` in PyTorch 1.10

[ananthsub]

Proposed refactoring or deprecation

Update the imports in the Lightning Quantization callback to torch.ao.quantization

Motivation

Quantization in PyTorch 1.10 is moving under the torch.ao module namespace: https://github.com/pytorch/pytorch/blob/3d6d4f4322e42886349b822449b9e439fac89ae2/torch/quantization/quantize.py#L2-L8

Lightning users on PyTorch 1.10+ shouldn't see any deprecation warnings because the framework imports from the old path.

Pitch

Update the imports here to reflect the PyTorch quantization import path changing: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/callbacks/quantization.py

We can gate this using https://github.com/PyTorchLightning/pytorch-lightning/blob/ab207921b9149a2c20ff9f797b6381f61cb09b1e/pytorch_lightning/utilities/imports.py#L74

BEFORE:

from torch.quantization import QConfig

AFTER:

if _TORCH_GREATER_EQUAL_1_10:
    import torch.ao.quantization as ptq
    from torch.ao.quantization import QConfig
else:
    import torch.quantization as ptq
    from torch.quantization import QConfig

Additional context

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, finetuning and solving problems with deep learning

• Bolts: Pretrained SOTA Deep Learning models, callbacks and more for research and production with PyTorch Lightning and PyTorch

• Lightning Transformers: Flexible interface for high performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

[theory-in-progress]

Hey @ananthsub , I'd like to work on this issue of refactoring the code from torch to torch.ao module namespace, could you please assign it to me and if possible could you provide some more context and elaborate on how to work on the issue? Thanks!

[ananthsub]

Hi @theory-in-progress , thanks for your interest! The pitch describes what we can do to safely transition to the new import path. Feel free to send a pull request and we'll be happy to review it 😄

[theory-in-progress]

Hey @ananthsub,, sorry for the delay. As I understand from the pitch, I have to update the imports in the file https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/callbacks/quantization.py from torch.quantization to torch.ao.quantization right? Please correct me if I'm wrong.

[ananthsub]

@theory-in-progress yes, the import path has to be updated if we're using PyTorch 1.10 or greater