mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
13ad4739a6
### Summary: This PR implements PTQ for APoT FakeQuant. It runs models (Resnet-18 pre-trained model, ImageNet dataset) to compare accuracy metrics for different qconfig settings of uniform vs. APoT quantized activation and weight. According to the collected accuracy stats, model #2 (uniform activation and APoT weight) appears to have a slight improvement in accuracy compared to model #1 (uniform activation and uniform weight) for 8-bit and significant improvement for 4-bit (see "Accuracy Stats" section below). ### Test Plan: Run models with: `python test/quantization/core/experimental/fx_graph_mode_apot.py` ### Accuracy Stats: 8-bit (Uniform int8, APoT b = 8 k = 2) **Model #1:** Uniform activation, uniform weight (FX Graph Mode quantized) Evaluation accuracy on test dataset: 64.43% (Top-1), 85.62% (Top-5) **Model #2:** Uniform activation, APoT weight (FX Graph Mode quantized) Evaluation accuracy on test dataset: 64.51% (Top-1), 85.78% (Top-5) **Model #3:** APoT activation, APoT weight (FX Graph Mode quantized) Evaluation accuracy on test dataset: 64.32% (Top-1), 85.78% (Top-5) 4-bit (Uniform int4, APoT b = 4 k = 2) **Model #1:** Uniform activation, uniform weight (FX Graph Mode quantized) Evaluation accuracy on test dataset: 45.63% (Top-1), 71.96% (Top-5) **Model #2:** Uniform activation, APoT weight (FX Graph Mode quantized) Evaluation accuracy on test dataset: 64.24% (Top-1), 85.56% (Top-5) **Model #3:** APoT activation, APoT weight (FX Graph Mode quantized) Evaluation accuracy on test dataset: 45.40% (Top-1), 76.21% (Top-5) **Full Precision model (FX Graph Mode quantized)** Evaluation accuracy on test dataset: 69.76% (Top-1), 89.08% (Top-5) **Eager mode quantized model** Evaluation accuracy on test dataset: 69.49% (Top-1), 88.90% (Top-5) Pull Request resolved: https://github.com/pytorch/pytorch/pull/81040 Approved by: https://github.com/jerryzh168
39 lines
1.6 KiB
Python
39 lines
1.6 KiB
Python
import torch
|
|
from torch import Tensor
|
|
from torch.ao.quantization.experimental.observer import APoTObserver
|
|
from torch.ao.quantization.fake_quantize import FakeQuantizeBase
|
|
from torch.ao.quantization.experimental.fake_quantize_function import fake_quantize_function
|
|
|
|
class APoTFakeQuantize(FakeQuantizeBase):
|
|
alpha: Tensor
|
|
gamma: Tensor
|
|
quantization_levels: Tensor
|
|
level_indices: Tensor
|
|
|
|
def __init__(self, observer=APoTObserver, **observer_kwargs):
|
|
super().__init__()
|
|
self.activation_post_process = observer(**observer_kwargs)
|
|
self.dtype = self.activation_post_process.dtype
|
|
|
|
def calculate_qparams(self, signed=False): # type: ignore[override]
|
|
return self.activation_post_process.calculate_qparams(signed=signed)
|
|
|
|
def forward(self, X: torch.Tensor): # type: ignore[override]
|
|
if self.observer_enabled[0] == 1:
|
|
self.activation_post_process.forward(X)
|
|
result = self.activation_post_process.calculate_qparams(signed=False)
|
|
self.alpha = result[0]
|
|
self.gamma = result[1]
|
|
self.quantization_levels = result[2]
|
|
self.level_indices = result[3]
|
|
|
|
if self.fake_quant_enabled[0] == 1:
|
|
assert (self.alpha is not None
|
|
and self.gamma is not None
|
|
and self.quantization_levels is not None
|
|
and self.level_indices is not None), "Must set qparams for fake quant"
|
|
|
|
X = fake_quantize_function.apply(X, self.alpha, self.gamma, self.quantization_levels, self.level_indices)
|
|
|
|
return X
|