Hello everyone,
I’m trying to perform Quantization-Aware Training (QAT) on YOLOv11 using Ultralytics. However, I’m facing an issue where my modifications to the model are being reset before training starts.
Here’s the code I wrote for it:
from ultralytics import YOLO
import torch
from torch.nn import Module
DEVICE = 0 if torch.cuda.is_available() else 'cpu'
class QAT(Module):
def __init__(self, model):
super().__init__()
self.model = model
self.quant = torch.quantization.QuantStub()
self.de_quant = torch.quantization.DeQuantStub()
self.yaml = model.yaml
def forward(self, x):
x = self.quant(x)
x = self.model(x)
for i in range(len(x)):
x[i] = self.de_quant(x[i])
return x
model = YOLO('yolo11n.pt')
qa_model = QAT(model.model)
qa_model.qconfig = torch.quantization.get_default_qconfig('qnnpack')
torch.quantization.prepare_qat(qa_model, inplace=True)
model.model = qa_model
model.train(
data='dataset/data.yaml',
epochs=100
)
It seems that Ultralytics resets the model to its default state before training starts, discarding my modifications.
To debug this, I added a callback:
def print_yolo_before_start(trainer):
print(trainer)
model.add_callback('on_train_start', print_yolo_before_start)
But when this runs, it prints the original YOLO model instead of the QAT-wrapped version, which differs in Conv layers.
How can I ensure my QAT modifications persist during training?