PTS-Quant

PyTorch implementation of PTS-Quant, a post-training quantization (PTQ) method that constrains quantization scales to powers-of-two for hardware-efficient neural network inference.

Overview

Quantization is an important technique for deploying deep neural networks on resource-constrained hardware.
By mapping floating-point weights and activations to low-bit integers, quantization can significantly reduce model size and computation cost.

However, conventional quantization methods typically use arbitrary floating-point scale factors, which still require floating-point multiplications during inference.

PTS-Quant addresses this issue by constraining quantization scales to power-of-two values, allowing multiplications to be replaced with efficient bit-shift operations on hardware.

The proposed method introduces a learnable scale rounding value and jointly optimizes:

Weight rounding
Scale rounding

within a rounding-based PTQ framework.

Key Features

Post-training quantization (PTQ)
Power-of-two scale quantization
Learnable scale rounding values
Joint optimization of weight and scale rounding values
Hardware-friendly quantization for efficient inference

Method

PTS-Quant extends rounding-based PTQ methods such as AdaRound and PD-Quant.

The quantization scale is constrained as:

scale = 2^k

where k is an integer learned through a scale rounding optimization process.

During reconstruction, PTS-Quant jointly optimizes:

Weight rounding values
Scale rounding values

This reduces quantization error while preserving the power-of-two constraint.

Quantization Pipeline

Full-Precision Model
        ↓
Calibration Dataset
        ↓
Scale and Weight Rounding Optimization
        ↓
Fix Scale to Power-of-two
        ↓
Fine-tune Weight Rounding Values
        ↓
Power-of-Two Quantized Model

Repository Structure

PTS-Quant
│
├── quant/
│   ├── ptq.py
│   ├── quant_layer.py
│   ├── quant_block.py
│   ├── quant_model.py
│   ├── fold_bn.py
│   ├── data_utils.py
│   ├── layer_recon.py
│   ├── set_act_quantize_params.py
│   ├── set_weight_quantize_params.py
│   └── block_recon.py
│
├── models/
│   ├── cifar10
│   │   ├── Resnet.py
│   │   ├── regnet.py
│   │   └── MobileNetV2.py
│   ├── imagenet
│   │   ├── Resnet.py
│   │   ├── regnet.py
│   │   └── MobileNetV2.py
│
├── configs/
│
├── utils/
│
├── result_csv/
│
├── docker.sh
│
└── README.md

Installation

First clone the repository with git clone ssh://git@gitlab.aislab.ee.ncku.edu.tw:3175/aislab-internal/research/active-research/model-compression/pts-quant.git

Create environment with docker

Build Docker Image

./docker.sh build

Run/Start Container

./docker.sh run

Dataset

ImageNet

Please download the ImageNet (ILSVRC2012) dataset and extract it to your local storage (recommended: HDD).

After extraction, organize the dataset into the following structure:

data/
└── ImageNet-1k/
        ├── train/
        │ ├── n01440764/
        │ ├── n01443537/
        │ └── ...
        └── val/
        │ ├── n01440764/
        │ ├── n01443537/
        │ └── ...

cifar10

Automatic download by torchvision

Note: Each class should be stored in a separate folder, which is required by standard PyTorch ImageFolder dataloaders.

Dataset Setup

The ImageNet dataset path is configured using symbolic links.

Create symbolic links to your local ImageNet training and validation datasets:

bash scripts/setup_imagenet.sh /path/to/train /path/to/val

Example:

bash scripts/setup_imagenet.sh \
    /storage/share/datasets/ImageNet-1k/ILSVRC/Data/CLS-LOC/train \
    /storage/share/datasets/ImageNet-1k/ILSVRC/Data/CLS-LOC/val

This creates:

data/ImageNet-1k/
├── train -> /path/to/train
└── val   -> /path/to/val

The expected dataset structure used by the project is:

data/ImageNet-1k/
├── train
└── val

Running Experiments

This project supports two modes:

reconstruction: perform PTQ calibration and save quantized checkpoints
evaluate: load a saved checkpoint and evaluate its performance

Reconstruction Example

Example: quantizing ResNet-18 Set the configuration in config/:

version: 0.1.0

models:
  - model_name: ResNet18
    dataset: imagenet
    mode: reconstruction
    save_name: ResNet18-i4sc
    wq_params: {'n_bits': 4, symmetric: False, 'channel_wise': True, 'scale_method': 'mse'}
    aq_params: {'n_bits': 4, symmetric: False, 'channel_wise': False, 'scale_method': 'mse',
                    'leaf_param': True, 'prob': 0.5}
    constraint_fn: 'sigmoid'   #constraint function for the rounding value
    initialization_fn: 'tanh'  #initialization function for the rounding value
    scale_iter: [0, 1000, 2500, 5000]
    joint_training: True
    result_path: result_csv/imagenet/ResNet18/stop_iteration/reconstruct.csv
    save_path: checkpoints/imagenet/ResNet18/stop_iteration/

Evaluate Example

The evaluate mode is used to verify the performance of a previously reconstructed quantized checkpoint without running PTQ reconstruction again.

During evaluation:

The quantized checkpoint is loaded from weight_path
Quantizers and rounding values are restored automatically
The model is switched to quantized inference mode
Top-1 / Top-5 accuracy is evaluated on the validation dataset

A saved checkpoint contains:

quantized model weights
quantizer states
scale rounding values
weight rounding values
experiment configuration

This allows users to directly reproduce and verify quantized model performance from saved checkpoints.

To evaluate a saved checkpoint:

  - model_name: ResNet18
    mode: evaluate
    save_name: ResNet18-i4sc
    result_path: result_csv/ResNet18/joint_training/evaluate.csv
    weight_path: checkpoints/ResNet18/joint_training/ResNet18-i4sc-joint_false_s2500.pth

Note

weight_path is required in evaluate mode
Reconstruction/calibration will NOT be performed
The checkpoint must be generated from reconstruction mode
Evaluation automatically enables:
- weight quantization
- activation quantization

Model

The following models are supported in imagenet:

ResNet18
ResNet50
MobileNetV2
RegNetX-600MF
RegNetX-3.2GF

The following models are supported in cifar10:

ResNet18
ResNet50
MobileNetV2
RegNetX-200MF
RegNetX-400MF

You can define multiple models in the config:

models:
  - model_name: ResNet18
    ...
  - model_name: ResNet50
    ...

Dataset

This code supports two datasets

imagenet
cifar10

You can define datasets in the config:

models:
  - model_name: ResNet18
    dataset: imagenet
    ...
  - model_name: ResNet18
    dataset: cifar10
    ...

Constraint function

Controls how the rounding value is mapped to a valid range during training:

sigmoid
tanh

Initialization function

Defines how rounding values are initialized:

sigmoid FP32-based initialization using inverse sigmoid
tanh FP32-based initialization using inverse tanh
zero Initialize to α = 0 → θ = 0.5
random Random initialization near α ≈ 0 → θ ≈ 0.5

Scale iteration (sacle_iter)

Defines the transition point between Stage 1 and Stage 2.

i < scale_iter → Stage 1 (scale optimization)
i ≥ scale_iter → Stage 2 (scale fixed to PoT) Supports multiple values for ablation:

scale_iter: [0, 1000, 2500, 5000]

Joint Training (joint_training)

Controls optimization behavior in Stage 1:

True Jointly update: - weight rounding value - scale rounding value
False Only update: - scale rounding value

Stage 2 always updates weight rounding only.

Result Path

Specifies the file path for saving experiment results.

Type: str
If not provided, the default path is result_csv/ImageNet.csv.
The parent directory will be automatically created if it does not exist.

Save Path

Specifies where to save quantized checkpoints (reconstruction mode only).

If a directory is provided, checkpoints will be saved as: {save_name}_s{scale_iter}.pth
If multiple scale_iter values are used, filenames are automatically adjusted to avoid overwriting.

<<<<<<< HEAD

Test Batch Size

Specifies the batch size used during evaluation.

Only used in evaluate mode
Default: 16

Example: test_batch_size: 32

=======

Weight path

The weight path you want to load your quantized checkpoints (evaluate mode only).

Type: str

Example Usage

After editing the configuration file, calibration can run: python quant/ptq.py --config config/W2A4/W2A4.yaml

evaluation can run: python quant/ptq.py --config config/W2A4/evaluate.yaml

Experimental Results

We evaluate PTS-Quant on ImageNet across multiple architectures under strict power-of-two (PoT) scale constraints. The proposed method is compared with both QAT-based approaches (e.g., TQT, HMQ) and reconstruction-based PTQ methods (e.g., PD-Quant).

Results show that PTS-Quant consistently improves accuracy over PD-Quant under PoT constraints, while narrowing the performance gap between PTQ and QAT methods.

4bit Quantization(W4A4)

ResNet-18

Method	Scheme	W/A	Top-1	Δ
HAWQv3	QAT	4/4	68.45	-3.02
RAPQ	PTQ	4/4	69.28	-1.80
PD-Quant (PoT)	PTQ	4/4	68.70	-2.31
PTS-Quant	PTQ	4/4	68.77	-2.24

ResNet-50

Method	Scheme	W/A	Top-1	Δ
TQT	QAT	4/8	74.40	-0.80
HMQ	QAT	3.55/8	76.30	+0.15
HAWQv3	QAT	4/4	74.24	-3.48
RAPQ	PTQ	4/4	74.64	-2.36
PD-Quant (PoT)	PTQ	4/4	74.03	-2.60
PTS-Quant	PTQ	4/4	74.59	-2.04

MobileNetV2

Method	Scheme	W/A	Top-1	Δ
HMQ	QAT	4.16/8	71.40	-0.48
RAPQ	PTQ	4/4	64.48	-8.01
PD-Quant (PoT)	PTQ	4/4	65.57	-7.05
PTS-Quant	PTQ	4/4	65.63	-6.99

RegNetX-600MF

Method	Scheme	W/A	Top-1	Δ
RAPQ	PTQ	4/4	69.59	-4.12
PD-Quant (PoT)	PTQ	4/4	69.33	-4.19
PTS-Quant	PTQ	4/4	69.68	-3.84

RegNetX-3.2GF

Method	Scheme	W/A	Top-1	Δ
RAPQ	PTQ	4/4	74.25	-4.11
PD-Quant (PoT)	PTQ	4/4	75.24	-3.22
PTS-Quant	PTQ	4/4	76.12	-2.34

2bit Quantization(W2A2)

ResNet-18

Method	Scheme	W/A	Top-1	Δ
PD-Quant (PoT)	PTQ	2/2	45.70	-25.31
PTS-Quant	PTQ	2/2	51.24	-19.77

ResNet-50

Method	Scheme	W/A	Top-1	Δ
HMQ	QAT	2.04/8	75.00	+0.15
PD-Quant (PoT)	PTQ	2/2	51.91	-24.72
PTS-Quant	PTQ	2/2	53.20	-23.43

MobileNetV2

Method	Scheme	W/A	Top-1	Δ
HMQ	QAT	2.22/8	65.70	-6.18
PD-Quant (PoT)	PTQ	2/2	3.43	-69.19
PTS-Quant	PTQ	2/2	5.25	-67.37

RegNetX-600MF

Method	Scheme	W/A	Top-1	Δ
PD-Quant (PoT)	PTQ	2/2	27.70	-45.82
PTS-Quant	PTQ	2/2	33.76	-39.76

RegNetX-3.2GF

Method	Scheme	W/A	Top-1	Δ
PD-Quant (PoT)	PTQ	2/2	42.00	-36.46
PTS-Quant	PTQ	2/2	49.17	-29.29

W2A4 Quantization(W2A4)

ResNet18

Method	Scheme	W/A	Top-1	Δ
RAPQ	PTQ	2/4	65.32	-5.76
PTS-Quant	PTQ	2/4	64.13	-6.88

ResNet50

Method	Scheme	W/A	Top-1	Δ
RAPQ	PTQ	2/4	69.71	-7.29
PTS-Quant	PTQ	2/4	69.12	-7.51

MobileNetV2

Method	Scheme	W/A	Top-1	Δ
RAPQ	PTQ	2/4	48.12	-24.37
PTS-Quant	PTQ	2/4	44.32	-28.30

RegNetX-600MF

Method	Scheme	W/A	Top-1	Δ
RAPQ	PTQ	2/4	61.48	-12.23
PTS-Quant	PTQ	2/4	61.91	-11.61

RegNetX-3.2GF

Method	Scheme	W/A	Top-1	Δ
RAPQ	PTQ	2/4	69.49	-8.87
PTS-Quant	PTQ	2/4	71.15	-7.31

Δ denotes the accuracy drop compared to the corresponding full-precision model.

Key Observations

PTS-Quant consistently outperforms PD-Quant under power-of-two constraints across all evaluated models.
The performance gap between PTQ and QAT is significantly reduced.
The improvement is especially notable in lightweight models such as RegNet.

Name		Name	Last commit message	Last commit date
config		config
models		models
quant		quant
result_csv		result_csv
scripts		scripts
utils		utils
.dockerignore		.dockerignore
.gitignore		.gitignore
Dockerfile		Dockerfile
README.md		README.md
docker.sh		docker.sh
requirements.txt		requirements.txt

Latest commit

History

Folders and files

Repository files navigation

PTS-Quant

Overview

Key Features

Method

Quantization Pipeline

Repository Structure

Installation

Create environment with docker

Dataset

ImageNet

cifar10

Dataset Setup

Running Experiments

Reconstruction Example

Evaluate Example

Note

Model

Dataset

Constraint function

Initialization function

Scale iteration (sacle_iter)

Joint Training (joint_training)

Result Path

Save Path

Test Batch Size

Weight path

Example Usage

Experimental Results

4bit Quantization(W4A4)

ResNet-18

ResNet-50

MobileNetV2

RegNetX-600MF

RegNetX-3.2GF

2bit Quantization(W2A2)

ResNet-18

ResNet-50

MobileNetV2

RegNetX-600MF

RegNetX-3.2GF

W2A4 Quantization(W2A4)

ResNet18

ResNet50

MobileNetV2

RegNetX-600MF

RegNetX-3.2GF

About

Topics

Resources

Stars

Watchers

Forks

Used by

Contributors

Languages