Improve the quality of quantized_qat models. (#640)
Summary:
Improves the quality of quantized model training:
- Added train_bs argument for easier tuning
- Keep train batch size consistent with the paper
Profile for mobilenet_v2_quantized_qat train (bs=96):

Profile for resnet50_quantized_qat train (bs=32):

Pull Request resolved: https://github.com/pytorch/benchmark/pull/640
Reviewed By: aaronenyeshi
Differential Revision: D33118989
Pulled By: xuzhao9
fbshipit-source-id: 52089e8e61b2642005bf4f8f3fb88bbbb7910280