Is it possible to perform quantization on densenet169 and how?

Question

I have been trying to performing quantization on a densenet model without success. I have been trying to implement pytorch post training static quantization. Pytorch has quantized versions of other models, but does not have for densenet. Is it possible to quantize the densenet architecture.

I have searched for tutorials on how to apply quantization on pre-trained models but i havent had any success.

Answer 1

Here's how to do this on DenseNet169 from torchvision:

from torch.ao.quantization import QuantStub, DeQuantStub
from torch import nn
from torchvision.models import densenet169, DenseNet169_Weights
from tqdm import tqdm
from torch.ao.quantization import HistogramObserver, PerChannelMinMaxObserver
import torch

# Wrap base model with quant/dequant stub
class QuantizedDenseNet169(nn.Module):
    def __init__(self):
        super().__init__()
        self.dn = densenet169(weights=DenseNet169_Weights.IMAGENET1K_V1)
        self.quant = QuantStub()
        self.dequant = DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.dn(x)
        return self.dequant(x)

dn = QuantizedDenseNet169()
# move to gpu
dn.cuda()

# Propagate qconfig
dn.qconfig = torch.quantization.QConfig(
    activation=HistogramObserver.with_args(),
    weight=PerChannelMinMaxObserver.with_args(dtype=torch.qint8)
)
# fbgemm for x86 architecture
torch.backends.quantized.engine = 'fbgemm'
dn = torch.quantization.prepare(dn, inplace=False)

# calibrate with own dataset (I'm using random inputs to show process)
with torch.no_grad():
    for _ in tqdm(range(5), desc="PTQ progess"):
        input_ = torch.randn([1, 3, 128, 128], device='cuda')
        dn.forward(input_)

# move to cpu before quantization
dn.cpu()
dn = torch.quantization.convert(dn, inplace=False)

# check if it's working
out = dn(torch.randn([1, 3, 128, 128]))