Weights not quantized after using qnn.QuantConv2d layers for QAT

[Maya7991]

I am using a custom CNN architecture and trained my model using qnn.QuantConv2d layers for QAT. When I inspect the stored weights of model, to load it to hardware it is not quantized. Given below is the model and weights,

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # First convolutional layer
        self.conv1 = qnn.QuantConv2d(in_channels=1, out_channels=16, kernel_size=3, padding=1, bias=True, weight_bit_width=4)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 =qnn.QuantConv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1, bias=True, weight_bit_width=4)
        self.conv3 = qnn.QuantConv2d(in_channels=32, out_channels=32, kernel_size=1, padding=0, bias=True, weight_bit_width=4)
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # First conv + ReLU + max pool
        x = self.pool(F.relu(self.conv2(x)))  # Second conv + ReLU + max pool
        x = F.relu(self.conv3(x))  # Third conv + ReLU + max pool
        x = x.view(-1, 32 * 7 * 7) 
        x = F.relu(self.fc1(x))     
        x = self.fc2(x)            
        return x

cnn_model = SimpleCNN()
cnn_loss=trainNet(cnn_model,0.0003)   # not including training loop here

state_dict = cnn_model.state_dict()
# Extract weights and biases only for the layer 'conv'
weights = {}
for name, param in state_dict.items():
  if 'conv.weight' in name:
        weights[name] = param.cpu().numpy()
        print(weights[name])

The weights that are printed looks like without quantization:

[[[[ 0.20773199 -0.17919816  0.01762806]
   [ 0.3506115  -0.15825677 -0.26748863]
   [-0.11684228  0.16537677  0.2590774 ]]]


 [[[ 0.19875231 -0.34671602  0.16409244]
   [-0.23938051 -0.84747297  0.23055074]
   [-0.37333268 -0.80596954 -0.17869644]]]
.....

If this is not the right way to access quantized weights, can you direct me how to access them?

Thanks,
Maya

[Giuseppe5]

Brevitas doesn't replace the original weights with quantized counterparts, instead you can generate them like this:

quant_weight = cnn_model.conv1.quant_weight()