Quantization aware training (QAT) with convolution output 32b quantization (for sanity check) accuracy drops too much - quantization

Question

I modified Distiller (https://github.com/IntelLabs/distiller) to emulate in-memory computing circuit, especially added a convolution layer quantization during QAT. However, accuracy drops over 60% (90% > 30%) even with 32b quantization for sanity check. I also want to say that it was perfectly fine when I just add a calculated quantization noise. Below is the code.
import torch
import torch.nn as nn
import math
__all__ = ['preact_resnet20_cifar', 'preact_resnet32_cifar', 'preact_resnet44_cifar', 'preact_resnet56_cifar',
'preact_resnet110_cifar', 'preact_resnet20_cifar_conv_ds', 'preact_resnet32_cifar_conv_ds',
'preact_resnet44_cifar_conv_ds', 'preact_resnet56_cifar_conv_ds', 'preact_resnet110_cifar_conv_ds']
NUM_CLASSES = 10
device = torch.device("cuda")
def conv3x3(in_planes, out_planes, stride=1):
"""3x3 convolution with padding"""
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
def quantize(tensor, n_bits=32, dim=144, clip_ratio=1):
max_value = dim*clip_ratio
min_value = -dim*clip_ratio
delta = max_value - min_value
y = tensor.clone()
y = torch.clamp(y, min = min_value, max = max_value)
lsb = delta / (2**n_bits)
y = (y // lsb)*lsb
return y
class PreactBasicBlock(nn.Module):
expansion = 1
def __init__(self, block_gates, inplanes, planes, stride=1, downsample=None, preact_downsample=True, resolution = 32):
super(PreactBasicBlock, self).__init__()
self.block_gates = block_gates
self.pre_bn = nn.BatchNorm2d(inplanes)
self.pre_relu = nn.ReLU(inplace=False) # To enable layer removal inplace must be False
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=False)
self.conv2 = conv3x3(planes, planes)
self.downsample = downsample
self.stride = stride
self.preact_downsample = preact_downsample
self.resolution = resolution
def forward(self, x):
need_preact = self.block_gates[0] or self.block_gates[1] or self.downsample and self.preact_downsample # add pre-activation for block_gates 0 and 1 and when is downsampled
if need_preact: # x > bn > relu > out
preact = self.pre_bn(x)
preact = self.pre_relu(preact)
out = preact
else: # x > out
preact = out = x
if self.block_gates[0]: # (preact) > conv > bn > relu
out = self.conv1(out)
dim = self.conv1.in_channels * self.conv1.kernel_size[0] * self.conv1.kernel_size[1]
out = quantize(out, n_bits=self.resolution, dim=dim, clip_ratio=1)
out = self.bn(out)
out = self.relu(out)
if self.block_gates[1]: # (preact)> conv
out = self.conv2(out)
dim = self.conv2.in_channels * self.conv2.kernel_size[0] * self.conv2.kernel_size[1]
out = quantize(out, n_bits=self.resolution, dim=dim, clip_ratio=1)
if self.downsample is not None:
if self.preact_downsample:
residual = self.downsample(preact)
else:
residual = self.downsample(x)
else:
residual = x
out += residual
return out
class PreactResNetCifar(nn.Module):
def __init__(self, block, layers, num_classes=NUM_CLASSES, conv_downsample=False):
self.nlayers = 0
# Each layer manages its own gates
self.layer_gates = []
for layer in range(3):
# For each of the 3 layers, create block gates: each block has two layers
self.layer_gates.append([]) # [True, True] * layers[layer])
for blk in range(layers[layer]):
self.layer_gates[layer].append([True, True])
self.inplanes = 16 # 64
super(PreactResNetCifar, self).__init__()
self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
self.layer1 = self._make_layer(self.layer_gates[0], block, 16, layers[0],
conv_downsample=conv_downsample)
self.layer2 = self._make_layer(self.layer_gates[1], block, 32, layers[1], stride=2,
conv_downsample=conv_downsample)
self.layer3 = self._make_layer(self.layer_gates[2], block, 64, layers[2], stride=2,
conv_downsample=conv_downsample)
self.final_bn = nn.BatchNorm2d(64 * block.expansion)
self.final_relu = nn.ReLU(inplace=True)
self.avgpool = nn.AvgPool2d(8, stride=1)
self.fc = nn.Linear(64 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def _make_layer(self, layer_gates, block, planes, blocks, stride=1, conv_downsample=False):
downsample = None
outplanes = planes * block.expansion
if stride != 1 or self.inplanes != outplanes:
if conv_downsample:
downsample = nn.Conv2d(self.inplanes, outplanes,
kernel_size=1, stride=stride, bias=False)
else:
# Identity downsample uses strided average pooling + padding instead of convolution
pad_amount = int(self.inplanes / 2)
downsample = nn.Sequential(
nn.AvgPool2d(2),
nn.ConstantPad3d((0, 0, 0, 0, pad_amount, pad_amount), 0)
)
layers = []
layers.append(block(layer_gates[0], self.inplanes, planes, stride, downsample, conv_downsample))
self.inplanes = outplanes
for i in range(1, blocks):
layers.append(block(layer_gates[i], self.inplanes, planes))
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.final_bn(x)
x = self.final_relu(x)
x = self.avgpool(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
def preact_resnet20_cifar(**kwargs):
model = PreactResNetCifar(PreactBasicBlock, [3, 3, 3], **kwargs)
return model
def preact_resnet32_cifar(**kwargs):
model = PreactResNetCifar(PreactBasicBlock, [5, 5, 5], **kwargs)
return model
def preact_resnet44_cifar(**kwargs):
model = PreactResNetCifar(PreactBasicBlock, [7, 7, 7], **kwargs)
return model
def preact_resnet56_cifar(**kwargs):
model = PreactResNetCifar(PreactBasicBlock, [9, 9, 9], **kwargs)
return model
def preact_resnet110_cifar(**kwargs):
model = PreactResNetCifar(PreactBasicBlock, [18, 18, 18], **kwargs)
return model
def preact_resnet182_cifar(**kwargs):
model = PreactResNetCifar(PreactBasicBlock, [30, 30, 30], **kwargs)
return model
def preact_resnet20_cifar_conv_ds(**kwargs):
return preact_resnet20_cifar(conv_downsample=True)
def preact_resnet32_cifar_conv_ds(**kwargs):
return preact_resnet32_cifar(conv_downsample=True)
def preact_resnet44_cifar_conv_ds(**kwargs):
return preact_resnet44_cifar(conv_downsample=True)
def preact_resnet56_cifar_conv_ds(**kwargs):
return preact_resnet56_cifar(conv_downsample=True)
def preact_resnet110_cifar_conv_ds(**kwargs):
return preact_resnet110_cifar(conv_downsample=True)
def preact_resnet182_cifar_conv_ds(**kwargs):
return preact_resnet182_cifar(conv_downsample=True)
I use distiller.example.classifier_compression.compress_classifier.py in Distiller, and use terminal with schedule file. Command is "python compress_classifier.py -a preact_resnet20_cifar --lr 0.1 -p 50 -b 128 -j 1 --epochs 200 --compress=../quantization/quant_aware_train/preact_resnet_cifar_dorefa.yaml --wd=0.0002 --vs=0 --gpus 0", for your reference.
I tried to add quantization noise, and accuracy was fine. But in my opinion it is not perfect emulation for quantization because data distribution is not uniform.