[Quant] Add FX support in quantization examples

[andrewor14]

Stack from ghstack (oldest at bottom):

• -> [Quant] Add FX support in quantization examples #5797

Summary: Previously, the quantization examples use only eager
mode quantization. This commit adds support for FX mode
quantization as well.

Test Plan:

# ==================== PTQ ====================
# MODEL is one of googlenet, inception_v3, resnet18, resnet50, resnext101_32x8d,
# shufflenet_v2_x0_5, shufflenet_v2_x1_0

# eager
python train_quantization.py --device="cpu" --post-training-quantize --backend="fbgemm"\
  --model="$MODEL" --weights="IMAGENET1K_V1" --quantization-workflow-type="eager_mode_quantization"

# fx
python train_quantization.py --device="cpu" --post-training-quantize --backend="fbgemm"\
  --model="$MODEL" --weights="IMAGENET1K_V1" --quantization-workflow-type="eager_mode_quantization"

# ==================== QAT ====================
# mobilenet_v2 eager
python train_quantization.py --device="cuda" --backend="qnnpack" --model="mobilenet_v2"\
  --epochs=10 --workers=64 --weights="IMAGENET1K_V1" --lr=0.0001 --weight-decay=0.0001\
  --quantization-workflow-type="eager_mode_quantization"

# mobilenet_v2 fx
python train_quantization.py --device="cuda" --backend="qnnpack" --model="mobilenet_v2"\
  --epochs=10 --workers=64 --weights="IMAGENET1K_V1" --lr=0.0001 --weight-decay=0.0001\
  --quantization-workflow-type="fx_graph_mode_quantization"

# mobilenet_v3_large eager
python train_quantization.py --device="cuda" --backend="qnnpack" --model="mobilenet_v3_large"\
  --epochs=10 --workers=64 --weights="IMAGENET1K_V1" --lr=0.001 --weight-decay=0.00001\
  --quantization-workflow-type="eager_mode_quantization"

# mobilenet_v3_large fx
python train_quantization.py --device="cuda" --backend="qnnpack" --model="mobilenet_v3_large"\
  --epochs=10 --workers=64 --weights="IMAGENET1K_V1" --lr=0.001 --weight-decay=0.00001\
  --quantization-workflow-type="fx_graph_mode_quantization"

Reviewers: jerryzh168, vkuzo

Subscribers: jerryzh168, vkuzo

[datumbox]

Thanks for the work @andrewor14. Some initial questions concerning the new API below.

[jerryzh168]

looks great! can you include the accuracy numbers in the summary and update the docs that shows accuracy as well

[andrewor14]

There was a problem with the way I set up the FX experiments. After fixing this problem I uncovered a new blocking issue that I summarized here pytorch/pytorch#75825. I will rerun all the experiments after we fix that issue.

[datumbox]

Thanks for the update. I've added some additional comments and questions. Let me know your thoughts.

[datumbox]

Just marking as "Request changes" to avoid accidental merges while we clarify some remaining questions.

[andrewor14]

Hi @datumbox, continuing our discussion here:

What steps if any do I have to take to load PTQ/QAT weights?

I'm trying to figure out what are the steps for loading pre-trained quantized weights to a model. Can this be done directly to the original model without any conversions? Or some FX convert call is needed prior loading?

For the test_only mode, we use quantized weights so convert must happen before loading the weights. For PTQ and QAT, however, we use pretrained FP32 weights so convert must happen after loading the weights. This is the same for eager mode and FX graph mode. In some more detail, here is the behavior today (pseudocode):

# PTQ / QAT eager
weights = get_pretrained_fp32_weights()
model = torchvision.models.quantization.QuantizableResNet50(...)
model.load_state_dict(weights.get_state_dict())
model = prepare(model) # or prepare_qat in QAT mode
calibrate(model, data) # or train in QAT mode
model = convert(model)
evaluate(model)

# test_only eager
model = torchvision.models.quantization.QuantizableResNet50(...)
model = prepare(model)
calibrate(model, dummy_data)
model = convert(model)
weights = get_quantized_weights()
model.load_state_dict(weights.get_state_dict())
evaluate(model)

What we want is the equivalent for FX. The only differences here are (1) We use prepare_fx (and the like) instead of prepare, and (2) We use the unmodified FP32 model (e.g. torchvision.models.ResNet50) instead of the model with manually inserted quant and dequant stubs (e.g. torchvision.models.quantization.QuantizableResNet50).

[datumbox]

Worth noting that we need to estimate these values with batch-size=1 on 1 GPU to avoid variances introduced due to batch padding (see #4559).

[datumbox]

@andrewor14 Thanks a lot for your work on this. Let's discuss the design doc you shared with me offline to decide on the approach and minimize throw away work. :)

[andrewor14]

As discussed offline, I'm closing this PR for now since the FX graph mode quantization API is subject to change. I will reopen this PR again and rerun all the experiments once the new API is ready. Thank you everyone for your comments so far!