Tối ưu hóa và Triển khai Mô hình YOLOv5 với TorchScript cùng TensorRT

Môi trường Cài Đặt & Yêu Cầu Hệ Thống

Để thực hiện huấn luyện và xuất mô hình ổn định, hệ thống nên đáp ứng các thành phần sau:
  • Hệ điều hành: Ubuntu 20.04 LTS
  • Ngôn ngữ: Python 3.8+
  • Deep Learning Framework: PyTorch >= 1.7.0
  • CUDA Toolkit: 11.3 trở lên
  • Multimedia Processing: OpenCV-Python 4.7+

Các gói phụ thuộc có thể được cài đặt tự động thông qua file yêu cầu của dự án nguồn.

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

Tổng Quan Kiến Trúc YOLOv5

YOLOv5 kế thừa những cải tiến từ YOLOv4 và chia kiến trúc thành bốn khối chính: tiền xử lý đầu vào, mạng trích xuất đặc trưng, mạch hợp nhất đa tỷ lệ và đầu ra dự đoán. Mỗi giai đoạn đều được tối ưu để cân bằng giữa tốc độ suy luận và độ chính xác.

Giai Đoạn Tiền Xử Lý Đầu Vào (Input)

Khối này chịu trách nhiệm chuẩn hóa hình ảnh trước khi đưa vào mạng. Các kỹ thuật nổi bật bao gồm:
  • Mosaic Augmentation: Ghép ngẫu nhiên 4 ảnh huấn luyện thành một khung hình duy nhất giúp mô hình học được ngữ cảnh không gian tốt hơn, đồng thời giảm bộ nhớ GPU trên mỗi batch.
  • Adaptive Anchor Calculation: Tự động tính toán kích thước và tỷ lệ khung đóng gói tối ưu dựa trên tập dữ liệu cụ thể thay vì dùng giá trị mặc định cố định.
  • Adaptive Image Scaling: Thay vì chỉ thêm viền đen đơn thuần, thuật toán tính toán tỷ lệ co giãn sao cho giữ nguyên tỉ lệ gốc đồng thời căn chỉnh theo stride của layer convolution và pooling, giảm thiểu padding thừa gây lãng phí tính toán.

Mạng Trích Xuất Đặc Trưng (Backbone)

Phần thân mạng kết hợp hai cấu trúc then chốt:
  • Focus Module: Hoạt động theo cơ chế slice 2x2 trên chiều kênh. Ví dụ, ảnh RGB 608x608 sẽ được tách thành ma trận 304x304 với 12 kênh sâu trước khi truyền qua layer convolution tiêu chuẩn. Cách tiếp cận này bảo toàn thông tin không gian mà không tốn chi phí pooling quá mức.
  • CSP1_X Block: Tích hợp tư tưởng Cross Stage Partial Network vào backbone, chia luồng gradient thành hai nhánh riêng biệt rồi fuse lại ở đầu ra. Giúp giảm số lượng phép tính冗余 và tăng khả năng lan truyền tín hiệu ngược.

Mạch Hợp Nhất Đa Tỷ Lệ (Neck)

YOLOv5 sử dụng kiến trúc FPN (Feature Pyramid Network) kết hợp PAN (Path Aggregation Network). Khác với phiên bản tiền nhiệm dùng convolution thường xuyên, YOLOv5 thay thế bằng CSP2_Bottleneck tại các nút fuse. Điều này nâng cao việc trộn đặc trưng mức độ nông và sâu, cải thiện đáng kể khả năng phát hiện vật thể nhỏ hoặc chồng lấp.

Đầu Ra & Hàm Mất mát (Head)

Phần detector cuối cùng áp dụng GIoU Loss thay cho CIoU/CIoU truyền thống, giúp hội tụ nhanh hơn khi các box dự đoán chưa chồng lên ground truth. Đối với bước lọc biên độ置信度 thấp, module post-processing tích hợp DIoU-NMS, sử dụng khoảng cách tâm để loại bỏ box trùng lặp hiệu quả hơn trước khi trả về kết quả cuối cùng.

Quy Trình Huấn Luyện & Chuẩn Bị Dữ Liệu Nhãn

Dưới đây là lệnh khởi chạy suy luận trên webcam, file ảnh, video hoặc thư mục trực tiếp:
python detect.py --weights yolov5s.pt --source 0  # webcam
python detect.py --weights yolov5s.pt --source cam_video.mp4
python detect.py --weights yolov5s.pt --source ./dataset_images/

Khi huấn luyện trên bộ dữ liệu COCO, người dùng chỉ cần thiết lập siêu tham số batch size phù hợp với VRAM:

python train.py --data coco.yaml --epochs 300 --weights '' --cfg yolov5n.yaml --batch-size 128
# yolov5s   --batch-size 64
# yolov5m   --batch-size 40
# yolov5l   --batch-size 24
# yolov5x   --batch-size 16

Tự Động Hóa Chuyển Đổi Từ Labelme Sang YOLO Format

Thay vì kịch bản thủ công cũ, mã dưới đây được viết lại theo hướng lập trình hàm, sử dụng `pathlib` và quản lý mapping lớp dinamic để đảm bảo tính mở rộng và rõ ràng:
import argparse
import json
from pathlib import Path
from typing import Dict, List, Tuple
from sklearn.model_selection import train_test_split

# Registry lưu ánh xạ tên label sang index
CLASS_REGISTRY: Dict[str, int] = {}

def get_class_index(cls_name: str) -> int:
    if cls_name not in CLASS_REGISTRY:
        CLASS_REGISTRY[cls_name] = len(CLASS_REGISTRY)
    return CLASS_REGISTRY[cls_name]

def calculate_normalized_bbox(width: float, height: float, xyxy: Tuple[float, float, float, float]) -> Tuple[float, float, float, float]:
    x_min, x_max, y_min, y_max = xyxy
    bbox_w = x_max - x_min
    bbox_h = y_max - y_min
    center_x = (x_min + x_max) / 2.0
    center_y = (y_min + y_max) / 2.0
    
    dw, dh = 1.0 / width, 1.0 / height
    return center_x * dw, center_y * dh, bbox_w * dw, bbox_h * dh

def transform_labelme_to_yolo(src_dir: str, dest_dir: str, val_ratio: float = 0.2) -> None:
    source = Path(src_dir)
    target_root = Path(dest_dir)
    
    json_files = sorted(source.glob("*.json"))
    if not json_files:
        raise FileNotFoundError(f"Không tìm thấy file .json trong {src_dir}")
        
    names = [f.stem for f in json_files]
    indices = list(range(len(names)))
    train_idx, val_idx = train_test_split(indices, test_size=val_ratio, random_state=42)
    
    subsets = {"train": train_idx, "valid": val_idx}
    
    for subset_name, idx_list in subsets.items():
        sub_dir = target_root / subset_name
        (sub_dir / "images").mkdir(parents=True, exist_ok=True)
        (sub_dir / "labels").mkdir(parents=True, exist_ok=True)
        
        for i in idx_list:
            basename = names[i]
            json_p = source / f"{basename}.json"
            img_p = source / f"{basename}.png"
            
            if not img_p.exists() or not json_p.exists():
                continue
                
            # Sao chép ảnh sang thư mục tương ứng
            (sub_dir / "images" / img_p.name).write_bytes(img_p.read_bytes())
            
            # Đọc metadata và convert nhãn
            with open(json_p, "r", encoding="utf-8") as f:
                meta = json.load(f)
                
            img_w, img_h = int(meta["imageWidth"]), int(meta["imageHeight"])
            label_lines = []
            
            for shape in meta["shapes"]:
                points = shape["points"]
                xs = [p[0] for p in points]
                ys = [p[1] for p in points]
                
                min_x, max_x = max(0, min(xs)), max(0, max(xs))
                min_y, max_y = max(0, min(ys)), max(0, max(ys))
                
                if max_x <= min_x or max_y <= min_y:
                    continue
                    
                norm_box = calculate_normalized_bbox(img_w, img_h, (min_x, max_x, min_y, max_y))
                cls_id = get_class_index(shape["label"])
                label_lines.append(f"{cls_id} {' '.join(map(str, norm_box))}")
                
            lbl_path = sub_dir / "labels" / f"{basename}.txt"
            lbl_path.write_text("\n".join(label_lines) + "\n" if label_lines else "")
            
    # Xuất cấu hình dataset YAML
    cfg_file = target_root / "data.yaml"
    ordered_classes = sorted(CLASS_REGISTRY.keys(), key=lambda k: CLASS_REGISTRY[k])
    
    cfg_data = {
        "train": "../train/images",
        "val": "../valid/images",
        "nc": len(ordered_classes),
        "names": ordered_classes
    }
    
    try:
        import yaml
        with open(cfg_file, "w", encoding="utf-8") as yf:
            yaml.dump(cfg_data, yf, default_flow_style=False, allow_unicode=True)
    except ImportError:
        # Fallback nếu chưa cài PyYAML
        with open(cfg_file, "w", encoding="utf-8") as yf:
            yf.write(f"train: ../train/images\nval: ../valid/images\nnc: {len(ordered_classes)}\nnames: {ordered_classes}\n")
            
    print(f"[DONE] Chuyển đổi hoàn tất. Total Classes: {len(ordered_classes)} | Target: {target_root}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Chuyển đổi nhãn Labelme sang định dạng YOLOv5")
    parser.add_argument("--input", type=str, required=True, help="Thư mục chứa file json gốc")
    parser.add_argument("--output", type=str, required=True, help="Thư mục đích xuất dataset")
    args = parser.parse_args()
    transform_labelme_to_yolo(args.input, args.output)

Chạy lệnh: `python converter_script.py --input ./raw_labelme/ --output ./yolo_dataset/`. Sau đó sao chép `data.yaml` vào thư mục `data/` của codebase YOLOv5 để bắt đầu huấn luyện.

Xuất Mô Hình Đưa Vào Sản Xuất

Xuất Định Dạng TorchScript

TorchScript hữu ích cho việc nhúng mô hình vào môi trường C++ hoặc triển khai trên thiết bị biên.
python export.py --weights weights/best.pt --include torchscript --device 0
Cơ chế nội tại sử dụng `torch.jit.trace` kèm theo việc serialize các metadata quan trọng như stride và danh sách class vào file config đính kèm:
def build_torchscript_checkpoint(model, dummy_input, save_path, enable_optimization):
    checkpoint_file = save_path.with_suffix(".ts")
    traced_model = torch.jit.trace(model, dummy_input, strict=False)
    
    meta_info = {
        "stride": int(max(m.stride for m in model.modules())),
        "names": model.names,
        "shape": tuple(dummy_input.shape)
    }
    extra_files = {"config.txt": json.dumps(meta_info)}
    
    target_fn = traced_model._save_for_lite_interpreter if enable_optimization else traced_model.save
    target_fn(checkpoint_file, _extra_files=extra_files)
    return checkpoint_file

Xuất Định Dạng TensorRT Engine

TensorRT mang lại lợi thế tối ưu hardware GPU chuyên biệt. Phiên bản 8.x trở đi hỗ trợ ONNX parser mạnh mẽ hơn.
python export.py --weights weights/best.pt --include engine --half --device 0
Đoạn mã xây dựng engine bên dưới chú trọng việc cấu hình optimization profile cho dynamic shape và kích hoạt FP16 nếu hardware hỗ trợ:
def compile_to_trt_engine(onnx_path, trt_save_path, workspace_size_gb=4, enable_fp16=False, is_dynamic=False):
    import tensorrt as trt
    from tensorrt.parsers import onnx_graphsurgeon as gs
    
    logger = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(logger)
    network_flags = int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network = builder.create_network(network_flags)
    
    parser = trt.OnnxParser(network, logger)
    if not parser.parse_from_file(str(onnx_path)):
        raise RuntimeError(f"Parse thất bại cho {onnx_path}")
        
    config = builder.create_builder_config()
    config.max_workspace_size = workspace_size_gb * (1 << 30)
    
    if is_dynamic:
        opt_profile = builder.create_optimization_profile()
        for inp in network.inputs:
            opt_profile.set_shape(inp.name, (1, *inp.shape[1:]), (4, *inp.shape[1:]), (8, *inp.shape[1:]))
        config.add_optimization_profile(opt_profile)
        
    if enable_fp16 and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
        
    with builder.build_engine(network, config) as engine:
        serialized_data = engine.serialize()
        trt_save_path.parent.mkdir(parents=True, exist_ok=True)
        trt_save_path.write_bytes(serialized_data)
        return trt_save_path

Khắc Phục Sự Cố Runtime Shared Library

Trong quá trình compile C++ interface hoặc chạy binary deploy, lỗi sau thường xuất hiện nếu đường dẫn thư viện chưa được hệ thống nhận diện:
error while loading shared libraries: libopencv_imgproc.so.XXX: cannot open shared object file
Cách xử lý chuẩn xác:
  1. Xác minh vị trí file `.so` còn tồn tại trên máy bằng lệnh `locate libopencv_imgproc.so.XXX` hoặc `find / -name "libopencv*.so*" 2>/dev/null`.
  2. Tạo hoặc chỉnh sửa file cấu hình linker tại `/etc/ld.so.conf.d/opencv.conf` và ghi đường dẫn chứa thư viện vào đó.
  3. Cập nhật cache hệ thống để linker tải lại danh sách thư viện chia sẻ: `sudo ldconfig`.
Ngoài ra, có thể kiểm tra trực tiếp dependency của binary đang chạy bằng `ldd `. Nếu còn dòng `=> not found`, thao tác cập nhật `ldconfig` sẽ khắc phục tức thì. Đảm bảo biến môi trường `LD_LIBRARY_PATH` cũng đã gán đúng đường dẫn trong profile shell nếu chạy ngầm qua cron hoặc systemd service.

Thẻ: PyTorch YOLOv5 tensorrt torchscript opencv

Đăng vào ngày 10 tháng 8 lúc 05:20