PHẦN 2: MASTERCLASS QAT PIPELINE – TUYỆT KỸ ĐẠT HIỆU NĂNG INT8 "VÔ HAO TỔN" CHO YOLOv8 / YOLO11

Lời mở đầu: Vượt qua giới hạn của PTQ

Trong đợt "xuất quân" trước, chúng ta đã cùng nhau giải mã cấu trúc đồ thị tính toán của YOLO và các bài toán lượng tử hóa căn bản. Nếu bạn vô tình lướt qua hoặc muốn ôn lại nội dung cốt lõi, hãy đọc ngay tại đây nhé:

🔗 Đọc lại bài viết trước: PHẦN 1: Tối ưu hóa "thần tốc" YOLO – Nghệ thuật ép cân từ FP32 xuống INT8 và giải mã cấu trúc đồ thị tính toán.

Thông thường, khi cần tăng tốc mô hình, cách tiếp cận nhanh nhất là PTQ (Post-Training Quantization) – lấy mô hình đã train xong, chạy vài chục batch dữ liệu mẫu để tính toán Scale/Zero-point (bước Calibrate) rồi mang đi deploy.

Tuy nhiên, với các mô hình object detection khắt khe như YOLO, PTQ thường làm sụt giảm từ 1% đến 3% mAP – một cái giá quá đắt trong môi trường production thực tế. Để giành lại từng phần trăm độ chính xác ấy, chúng ta phải nâng cấp lên vũ khí tối tân hơn: QAT (Quantization-Aware Training).

Bài viết này sẽ bóc tách chi tiết pipeline 5 bước xây dựng hệ thống QAT hoàn chỉnh dựa trên thư viện pytorch_quantization của NVIDIA.

1. Bản chất cốt lõi: Tại sao phải so sánh QAT với PTQ?

Trước khi bắt tay vào gõ code, bạn cần hiểu rõ sơ đồ tư duy chiến lược và sự khác biệt bản chất giữa hai phương pháp này:

              ┌─────────────────────────────────────────┐
              │  Replace modules + Insert QDQ (Giống)   │
              │  Calibrate histogram + Compute amax     │
              └──────────────┬──────────────────────────┘
                             │
              ┌──────────────┴──────────────┐
              │                             │
         ┌────▼────┐                  ┌─────▼─────┐
         │   PTQ   │                  │    QAT    │
         │  STOP   │                  │ Fine-tune │
         │  here   │                  │ với MSE   │
         │         │                  │ teacher   │
         └────┬────┘                  └─────┬─────┘
              │                             │
         Weights giữ                  Weights được
         nguyên từ FP32              điều chỉnh để bù
                                     quantization error

Mẹo thực chiến: Luôn luôn chạy đánh giá (Eval) mô hình PTQ trước khi làm QAT. Nếu bước PTQ chỉ làm giảm một lượng cực nhỏ mAP (ví dụ ~0.2%), bạn hoàn toàn có thể đem deploy ngay lập tức mà không cần tốn tài nguyên chạy QAT làm gì cho mệt!

2. Giải mã Chi tiết Toàn bộ QAT Pipeline 5 Bước

Quy trình QAT thực tế trong mã nguồn được triển khai một cách nghiêm ngặt qua 5 bước tuần tự sau:

Bước 1: Initialize (Khởi tạo Descriptor mặc định)

Hệ thống gọi lệnh quantize.initialize() để thiết lập bộ cấu hình mặc định (Histogram Calibration) cho toàn bộ các module lượng tử hóa (QuantConv2d, QuantLinear,...). Bộ cấu hình này có nhiệm vụ chuẩn bị sẵn các "bộ nhớ" để thu thập phân phối dữ liệu kích hoạt ở các bước sau.

def initialize():
    quant_desc_input = QuantDescriptor(calib_method="histogram")
    quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
    quant_nn.QuantMaxPool2d.set_default_quant_desc_input(quant_desc_input)
    quant_nn.QuantAvgPool2d.set_default_quant_desc_input(quant_desc_input)
    quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
    quant_logging.set_verbosity(quant_logging.ERROR)

Bước 2: Replace Modules (Chèn node Fake-Quantize vào Model)

Bước này chia làm hai nhánh để biến đổi hoàn toàn đồ thị tính toán gốc sang đồ thị lượng tử hóa:

• 2a. Thay thế layer chuẩn: Quét toàn bộ mô hình và ép các lớp nn.Conv2d thành QuantConv2d. Mỗi mô hình lúc này sẽ tự động sinh ra các cặp node Q/DQ (Quantize/Dequantize) ngay trước Weight và Input của lớp đó.
• 2b. Thay thế các toán tử tùy biến (Custom Operations): Những phép toán như cộng tensor (+) không tự lượng tử hóa được. Chúng ta phải chủ động xây dựng module thay thế (như lớp QuantAdd) để ép dòng chảy dữ liệu đi qua TensorQuantizer.

class QuantAdd(torch.nn.Module):
    """Quantized Add operation for YOLOv8"""
    def __init__(self, quantization):
        super().__init__()
        if quantization:
            self._input0_quantizer = quant_nn.TensorQuantizer(QuantDescriptor(num_bits=8, calib_method="histogram"))
            self._input1_quantizer = quant_nn.TensorQuantizer(QuantDescriptor(num_bits=8, calib_method="histogram"))
            self._input0_quantizer._calibrator._torch_hist = True
            self._input1_quantizer._calibrator._torch_hist = True
            self._fake_quant = True
        self.quantization = quantization

    def forward(self, x, y):
        if self.quantization:
            return self._input0_quantizer(x) + self._input1_quantizer(y)
        return x + y

Sau đó, áp dụng kỹ thuật khỉ đột (monkey-patching) hoặc ghi đè thuộc tính để đổi hàm forward của khối Bottleneck nguyên bản:

def replace_bottleneck_forward_yolov8(model):
    """Replace Bottleneck forward method with quantized version for YOLOv8"""
    for name, bottleneck in model.named_modules():
        if bottleneck.__class__.__name__ == "Bottleneck":
            if hasattr(bottleneck, 'add') and bottleneck.add:
                if not hasattr(bottleneck, "addop"):
                    print(f"Add QuantAdd to {name}")
                    bottleneck.addop = QuantAdd(True)
                bottleneck.__class__.forward = bottleneck_quant_forward_yolov8

Bước 3: Calibrate (Thu thập thống kê để tính toán Scale)

Mô hình kích hoạt chế độ "tắt fake-quant, bật thu thập". Ta đẩy khoảng 25 batch dữ liệu (chỉ cần dữ liệu thô không nhãn) đi qua mô hình thông qua hàm collect_stats. Hệ thống sử dụng thuật toán tối ưu hóa MSE (Mean Squared Error) để tự động tìm ra điểm cắt tối ưu ($a_{\max}$), loại bỏ các giá trị nhiễu (outliers) quá lớn.

# Trích đoạn logic chuyển đổi trạng thái trong hàm collect_stats
for name, module in model.named_modules():
    if isinstance(module, quant_nn.TensorQuantizer):
        if module._calibrator is not None:
            module.disable_quant() # Tắt fake-quant để lấy phân phối chuẩn float
            module.enable_calib() # Kích hoạt bộ thu thập thống kê

Sau khi thu thập xong dữ liệu phân phối đồ thị, ta gọi cấu hình phương pháp MSE để chốt hạ bài toán tìm điểm bão hòa dữ liệu: compute_amax(model, device, method="mse").

Bước 4: Apply Custom Rules (Đảm bảo tính nhất quán - Quantizer Consistency)

Khi các luồng dữ liệu hội tụ tại phép Concat hoặc Add, nếu mỗi nhánh sở hữu một giá trị Scale khác nhau, TensorRT sẽ bắt buộc phải chèn thêm các toán tử Re-quantize để đồng bộ hóa dải dữ liệu. Điều này làm tăng chi phí phần cứng (overhead) và giảm FPS trầm trọng.

Giải pháp: Duyệt cấu trúc đồ thị tính toán để tìm ra các cặp phần tử phụ thuộc nhau, sau đó ép chúng chia sẻ chung một đối tượng cấu hình Python (_input_quantizer).

# Trích đoạn ép quy tắc đồng bộ hóa cho khối Bottleneck Skip-Connection
for name, bottleneck in model.named_modules():
    if bottleneck.__class__.__name__ == "Bottleneck":
        if hasattr(bottleneck, 'add') and bottleneck.add and hasattr(bottleneck, 'addop'):
            print(f"Rules: {name}.add match to {name}.cv1")
            if hasattr(bottleneck.cv1, 'conv'):
                major = bottleneck.cv1.conv._input_quantizer
            else:
                major = bottleneck.cv1._input_quantizer
            # Ép 2 nhánh đầu vào của phép cộng dùng chung vùng tham chiếu con trỏ quantizer
            bottleneck.addop._input0_quantizer = major
            bottleneck.addop._input1_quantizer = major

Bước 5: Fine-tune (Layer-wise Distillation với FP32 Teacher)

Đây chính là "ma thuật" giúp phục hồi độ chính xác gần như nguyên vẹn. Ta áp dụng cơ chế Chưng cất tri thức (Knowledge Distillation) cấp độ Layer bằng cách so sánh đầu ra từng lớp giữa mô hình gốc (FP32 Teacher) và mô hình đang cấu hình (INT8 Student) thông qua hàm torch.nn.MSELoss().

# Logic cốt lõi của vòng lặp Fine-tune chưng cất tri thức cấp độ layer
with amp.autocast(enabled=fp16):
    model(imgs) # Forward pass trên INT8 Student (có Fake-Quantize)

    with torch.no_grad():
        origin_model(imgs) # Forward pass trên FP32 Teacher (đã freeze gradient)

    quant_loss = 0
    # Tính tổng lỗi tích lũy giữa các cặp layer được giám sát
    for index, (mo, fo) in enumerate(zip(model_outputs, origin_outputs)):
        if isinstance(mo, torch.Tensor) and isinstance(fo, torch.Tensor):
            if mo.shape == fo.shape:
                quant_loss += quant_lossfn(mo, fo)

# Cập nhật trọng số thông qua kỹ thuật STE tích hợp sẵn
scaler.scale(quant_loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()

3. Tổng kết Toàn bộ Luồng Đi (E2E Flow Chart)

Để dễ hình dung khi triển khai thực tế trên hệ thống Notion/Viblo của bạn, hãy lưu lại sơ đồ chuỗi hành động này:

1. initialize() $\rightarrow$ Thiết lập histogram mặc định.
2. replace_modules() $\rightarrow$ Ép Conv thành QuantConv, tích hợp các cấu trúc QuantAdd, QuantConcat độc quyền.
3. calibrate_model() $\rightarrow$ Chạy dữ liệu mẫu $\rightarrow$ Vẽ Histogram $\rightarrow$ Tối ưu MSE để chốt giá trị Scale.
4. apply_custom_rules() $\rightarrow$ Quét đồ thị ONNX, ép các nhánh Concat/Skip-connection dùng chung một con trỏ Quantizer object.
5. finetune() $\rightarrow$ Ép Student INT8 học theo Teacher FP32 bằng Layer-wise MSE Loss qua 10 Epochs.
6. export ONNX $\rightarrow$ Đóng gói kèm plugin cấu hình $\rightarrow$ Sẵn sàng biên dịch sang TensorRT Engine chạy mượt mà trên phần cứng.

Lời kết

Lượng tử hóa nhận thức huấn luyện (QAT) không đơn thuần là một kỹ thuật viết mã, nó là một nghệ thuật giao thoa giữa toán học tối ưu phân phối dữ liệu và cấu trúc phần cứng cứng nhắc. Làm chủ được QAT Pipeline cho các dòng mạng phức tạp như YOLOv8 hay YOLO11 sẽ giúp bạn tự tin giải quyết những bài toán triển khai AI khó nhằn nhất trong môi trường doanh nghiệp.

🛠️ Kho tài nguyên thực chiến: Bạn có thể tham khảo đầy đủ mã nguồn triển khai, các quy tắc tối ưu đồ thị nâng cao (Custom Rules) và toàn bộ mã lệnh thử nghiệm tại kho lưu trữ GitHub chính thức: thanhtungdo2211/optimize-dl-model.

Nếu thấy bài viết hữu ích, đừng quên để lại một Upvote và Bookmark lại bài viết này trên Viblo của bạn nhé!