Quantization aware training（QAT）10x~slower

[Jon-drugstore]

📚 The doc issue

I have test the QAT using the edgeai-torchvision/references/edgeailite/scripts/train_classification_main.py for float, quantize, Freezing BN, Freezing ranges training, the results show that it is upto 10x slower than without the quantization, it is extremely slow-on gpu. The detail logs is following,

• Float Train

=> training 0.00% of 1x98...Epoch=2/10 LR=0.02000 Time=0.823 DataTime=0.665 Loss=2.270 Prec@1=14.844 Prec@5=60.156 rate=0 Hz, eta=?, total=0:00:00 => training 1.02% of 1x98...Epoch=2/10 LR=0.02000 Time=0.823 DataTime=0.665 Loss=2.270 Prec@1=14.844 Prec@5=60.156 rate=1522.87 Hz, eta=0:00:00, total=0:00:00 => training 1.02% of 1x98...Epoch=2/10 LR=0.02000 Time=0.823 DataTime=0.665 Loss=2.270 Prec@1=14.844 Prec@5=60.156 rate=1522.87 Hz, eta=0:00:00, total=0:00:00

Float training time: Time - DataTime ~= 160ms

• Quantize Train

=> training_quantize 0.00% of 1x98...Epoch=4/10 LR=0.00001 Time=2.107 DataTime=0.890 Loss=0.628 Prec@1=77.344 Prec@5=98.828 rate=0 Hz, eta=?, total=0:00:00 => training_quantize 1.02% of 1x98...Epoch=4/10 LR=0.00001 Time=2.107 DataTime=0.890 Loss=0.628 Prec@1=77.344 Prec@5=98.828 rate=1349.39 Hz, eta=0:00:00, total=0:00:00 => training_quantize 1.02% of 1x98...Epoch=4/10 LR=0.00001 Time=2.107 DataTime=0.890 Loss=0.628 Prec@1=77.344 Prec@5=98.828 rate=1349.39 Hz, eta=0:00:00, total=0:00: Quantize training time: Time - DataTime ~= 1200ms

• Freezing BN

=> training_quantize 0.00% of 1x98...Epoch=6/10 LR=0.00001 Time=2.002 DataTime=0.883 Loss=0.637 Prec@1=79.492 Prec@5=98.438 rate=0 Hz, eta=?, total=0:00:00 => training_quantize 1.02% of 1x98...Epoch=6/10 LR=0.00001 Time=2.002 DataTime=0.883 Loss=0.637 Prec@1=79.492 Prec@5=98.438 rate=1171.45 Hz, eta=0:00:00, total=0:00:00 => training_quantize 1.02% of 1x98...Epoch=6/10 LR=0.00001 Time=2.002 DataTime=0.883 Loss=0.637 Prec@1=79.492 Prec@5=98.438 rate=1171.45 Hz, eta=0:00:00, total=0:00:00 Freezing BN Quantize training time: Time - DataTime ~= 1200ms

• Freezing ranges

=> training_quantize 0.00% of 1x98...Epoch=8/10 LR=0.00000 Time=1.498 DataTime=0.926 Loss=0.707 Prec@1=74.023 Prec@5=98.242 rate=0 Hz, eta=?, total=0:00:00 => training_quantize 1.02% of 1x98...Epoch=8/10 LR=0.00000 Time=1.498 DataTime=0.926 Loss=0.707 Prec@1=74.023 Prec@5=98.242 rate=1176.00 Hz, eta=0:00:00, total=0:00:00 => training_quantize 1.02% of 1x98...Epoch=8/10 LR=0.00000 Time=1.498 DataTime=0.926 Loss=0.707 Prec@1=74.023 Prec@5=98.242 rate=1176.00 Hz, eta=0:00:00, total=0:00:00 Freezing ranges Quantize training time: Time - DataTime ~= 570ms

The result is expected or I did something wrong? Thanks a lot.

Suggest a potential alternative/fix

No response