I believe the issue is due to the insertion of observers. I’m currently running MobileNetV2 with QAT, but only evaluating with calibration batches (no actual training), and then performing inference (after freezing the observers).
When using observer=torch.quantization.MinMaxObserver, before running qat_model.apply(torch.quantization.disable_observer), the time per batch is 2.95s, and afterwards, it is 1.51s, on cpu. On a GPU, the numbers are 3.71s and 0.38s. When the observers are not enabled, there is a clear speedup on GPUs, but there’s actually a performance degradation with observers on CPU.
When using the HistogramObserver, the performance on GPU (just for calibrating, not training) is 8x worse.