I am working on a Yolo class object detection model and was able to do QAT in Static eager mode. With Pytorch 2.5 I am trying Pytorch 2 Export dynamic QAT, the code runs without any error, but loss and eval mAP don’t change with training epochs.
Here is my QAT preparation code. (Code is running on CPU). Once I receive the PT2E model from the below code training is done as usual.
from torch.ao.quantization.quantizer.xnnpack_quantizer import XNNPACKQuantizer, get_symmetric_quantization_config
from torch.ao.quantization.quantize_pt2e import (
prepare_qat_pt2e,
convert_pt2e,
)
from torch.export import Dim
batch_size = Dim("batch", min=1, max=512) # Batch size between 1 and 512
_height = Dim("_height", min=2) # Base height
_width = Dim("_width", min=2) # Base width
gt = Dim("gt", min=1, max=120) # Ground truth dimension
# Adjusted dynamic dimensions (height and width are multiples of 32)
height = 32 * _height
width = 32 * _width
example_inputs = (
torch.rand(2, 3, test_size[0], test_size[1]),
# torch.rand(2, 120, 5)
)
dynamic_shapes=(
{0: batch_size, 2: height, 3: width},
# {0: batch_size, 1:gt }
)
float_model = deepcopy(model)
exported_model = torch.export.export_for_training(
float_model.to('cpu'),
example_inputs,
dynamic_shapes=dynamic_shapes
).module()
quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())
prepared_model = prepare_qat_pt2e(exported_model, quantizer)
original_weights = torch.load(ckpt_file)['model']
prepared_model.load_state_dict(original_weights, strict=False)
return prepared_model
Am I missing something?