It seems disable_fake_quant and other similar functions not working for the following case (Pytorch 1.7).
class LeNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(28 * 28, 10)    
        self.relu1 = nn.ReLU(inplace=True)

    def forward(self, x):             
        return self.relu1(self.l1(x.view(x.size(0), -1)))  

fq_weight = torch.quantization.FakeQuantize.with_args(\
    observer=torch.quantization.MovingAverageMinMaxObserver.with_args(),
    quant_min=0, quant_max=255, dtype=torch.quint8)   

fq_activation = torch.quantization.FakeQuantize.with_args(\
    observer=torch.quantization.MovingAverageMinMaxObserver.with_args(),
    quant_min=0, quant_max=255, dtype=torch.quint8)   


model = LeNet()
model.l1.qconfig = torch.quantization.QConfig(activation=fq_activation, weight=fq_weight)     
torch.quantization.prepare_qat(model, inplace=True)
model.l1.apply(torch.quantization.disable_fake_quant)

image1011×265 21.1 KB
The above case shows it works just fine. Now, I make my own FakeQuantize as follows, then it doesn’t work any more.
class MyFakeQuantize(torch.quantization.FakeQuantize):
    def __init__(self, observer, quant_min, quant_max, n_cluster=0, **observer_kwargs):
        super().__init__(observer, quant_min, quant_max, **observer_kwargs)

fq_weight = MyFakeQuantize.with_args(\
    observer=torch.quantization.MovingAverageMinMaxObserver.with_args(),
    quant_min=0, quant_max=255, dtype=torch.quint8)   

fq_activation = MyFakeQuantize.with_args(\
    observer=torch.quantization.MovingAverageMinMaxObserver.with_args(),
    quant_min=0, quant_max=255, dtype=torch.quint8)   


model2 = LeNet()
model2.l1.qconfig = torch.quantization.QConfig(activation=fq_activation, weight=fq_weight)     
torch.quantization.prepare_qat(model2, inplace=True)
model2.l1.apply(torch.quantization.disable_fake_quant)

image1009×259 21.2 KB
Is this expected behavior?