How to skip quantizing specific layers in fx quantization QAT

Hi there,

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Sequential(
                nn.Conv2d(3, 128, 3, 1, 1), 
                nn.BatchNorm2d(128),
                nn.ReLU(inplace=True),
                nn.Conv2d(128, 128, 3, 1, 1), 
                nn.BatchNorm2d(128),
                nn.ReLU(inplace=True),
        )
    
    def forward(self, x): 
        return self.conv(x)


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

        self.backbone = Net()
    
        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) 
    
        self.fc = nn.Linear(128, 32) 
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x): 
        x = self.backbone(x)

        x = self.avg_pool(x)
        x = x.view(x.size(0), -1) 
        x = self.fc(x)
        x = self.relu(x)
    
        return x

If I want to attach observers (weight+activation) to backbone layers only, and keep the rest layers (self.avg_pool, self.fc, self.relu) unquantized, how should I define the qconfig_dict?

1 Like