Xuất model Ultralytics YOLO sang Ambarella CVflow#

Việc triển khai các model Ultralytics YOLO trên SoC Ambarella đòi hỏi phải biên dịch model bằng công cụ biên dịch của Ambarella; các model được tối ưu cho kiến trúc CVflow có hiệu năng suy luận tốt hơn. Bản fork này của Ultralytics tích hợp trực tiếp bộ công cụ nén SpongeTorch của Ambarella vào pipeline huấn luyện, xác thực và xuất model, giúp developer tạo ra các model được tối ưu để triển khai hiệu quả trên phần cứng Ambarella.

Hướng dẫn này trình bày quy trình triển khai phát hiện đối tượng hiện tại, từ huấn luyện có xét đến nén đến suy luận trên thiết bị (xem Tổng quan quy trình để biết đầy đủ các bước).

Định dạng checkpoint AmbaPB là phần mở rộng đặc thù của Ambarella cho đặc tả ONNX IR, hỗ trợ các primitive tính toán CVflow và đóng gói các artifact do công cụ SDK tạo ra. Đây là artifact phía host được dùng để xác thực độ chính xác của model đã biên dịch trước khi triển khai; binary Cavalry riêng biệt được tạo cho thiết bị đích mới là thành phần chạy trên board.

Ambarella là công ty nào?#

Ambarella, có trụ sở chính tại Santa Clara, California, là công ty bán dẫn chuyên thiết kế SoC AI biên. Bộ xử lý của công ty kết hợp xử lý tín hiệu hình ảnh, mã hóa video và năng lực tính toán AI trên chip, được sử dụng trong các thiết bị an ninh, ô tô, robot, công nghiệp và tiêu dùng.

CVflow là gì?#

CVflow là kiến trúc xử lý thị giác của Ambarella. Kiến trúc này sử dụng một engine xử lý thị giác chuyên dụng, tách biệt với CPU và GPU, để chạy các workload thị giác máy tính và mạng nơ-ron. Các model được huấn luyện bằng những framework như PyTorch sẽ được biên dịch sang định dạng gốc của CVflow bằng SDK Ambarella trước khi chạy trên engine.

Các dòng SoC CVflow hiện tại và ứng dụng phổ biến:

Tại sao nên triển khai YOLO trên Ambarella?#

• Hiệu năng trên mỗi watt: SoC CVflow được thiết kế cho AI biên hoạt động liên tục, chạy phát hiện đối tượng theo thời gian thực trong giới hạn công suất của camera.
• Huấn luyện có xét đến nén: SpongeTorch áp dụng pruning trong quá trình huấn luyện để giúp model duy trì độ chính xác, đồng thời trở nên thưa hơn và hiệu quả hơn khi triển khai trên CVflow.
• Pipeline camera tích hợp: SoC Ambarella kết hợp bộ xử lý tín hiệu hình ảnh (ISP), mã hóa video siêu HD và CVflow để hỗ trợ nhiều hệ thống camera với mức tiêu thụ điện năng thấp; nhờ đó, một SoC Ambarella duy nhất có thể xử lý toàn bộ pipeline camera AI.

Tổng quan quy trình#

Pipeline gồm sáu giai đoạn:

1. Huấn luyện có xét đến nén — huấn luyện với cấu hình SpongeKit (amba_config) để SpongeTorch áp dụng pruning phi cấu trúc dần dần trong quá trình huấn luyện. Khi độ chính xác của lượng tử hóa sau huấn luyện (PTQ) không đạt yêu cầu, SpongeTorch cũng hỗ trợ huấn luyện có xét đến lượng tử hóa (QAT), nhưng hiện quy trình này chưa được tích hợp vào bản tích hợp Ultralytics và dự kiến sẽ có trong bản phát hành tương lai.
2. Xuất ONNX — xuất checkpoint đã nén bằng cùng amba_config, giữ nguyên cấu trúc nén trong đồ thị ONNX.
3. Biên dịch — biên dịch model ONNX thành checkpoint AmbaPB bằng công cụ biên dịch SDK; các công cụ này áp dụng PTQ cho engine CVflow.
4. Xác thực trên host — chạy model *.ambapb.ckpt.onnx đã biên dịch thông qua Ultralytics predict/val bằng backend AmbaPB để xác minh độ chính xác trước khi triển khai.
5. Chuyển đổi Cavalry — chuyển checkpoint AmbaPB đã xác thực thành binary Cavalry bằng công cụ SDK.
6. Chạy trên thiết bị — chạy binary Cavalry trên thiết bị bằng thư viện runtime SDK của Ambarella.

Quy trình huấn luyện và xuất model bằng SpongeTorch là tùy chọn và có thể thay bằng cách xuất ONNX thông thường (xem Xuất model không dùng SpongeTorch).

Điều kiện tiên quyết#

Cài đặt#

Cài đặt bản fork Ultralytics này, sau đó thiết lập SDK Ambarella CVflow — bao gồm công cụ biên dịch và thư viện cvflowbackend — rồi cài đặt wheel spongetorch được phân phối kèm theo:

# Install this Ultralytics fork from source
git clone https://github.com/Ambarella-Inc/ultralytics
cd ultralytics
git checkout amba_v8.4.46
pip install -e .

# Access and set up the Ambarella SDK compilation tools
# After the environment is ready, install the spongetorch library
pip install /path/to/spongetorch-*.whl

AutoBackend tìm cvflowbackend thông qua lệnh tv2 (tv2 -libpath cvflowbackend) của công cụ biên dịch SDK; vì vậy, cần cài đặt công cụ biên dịch SDK và thêm công cụ này vào PATH trước khi chạy suy luận hoặc xác thực bằng model đã biên dịch.

File cấu hình SpongeKit#

SpongeTorch sử dụng file cấu hình SpongeKit (định dạng protobuf-text, .prototxt) để xác định các lượt pruning, bao gồm mục tiêu độ thưa và lịch trình nén. Lấy các cấu hình mẫu và tài liệu schema tương ứng từ bản phát hành SDK Ambarella của bạn. Để đảm bảo tính nhất quán giữa huấn luyện, xác thực và triển khai, hãy dùng cấu hình huấn luyện bất cứ khi nào quá trình xác thực cần chuẩn bị lại model, đồng thời luôn dùng cùng cấu hình khi xuất checkpoint đã nén.

Các tham số Amba#

Hai tham số kiểm soát tích hợp SpongeTorch trong các chế độ train, val và export:

Bản fork cũng bổ sung một tham số xuất model tổng quát:

Huấn luyện có xét đến nén#

Huấn luyện (hoặc fine-tune) model khi bật tính năng nén SpongeTorch:

from ultralytics import YOLO

model = YOLO("yolo26n.pt")
model.train(
    data="coco8.yaml",
    epochs=100,
    amba_config="config.prototxt",
    amba_chipset="CV72",
)

Khi đặt amba_config, trainer sẽ bọc model và optimizer bằng spongetorch.prepare() trong quá trình thiết lập. Nén được áp dụng dần theo lịch trình từng bước, giúp mạng học cách duy trì độ chính xác trong khi trở nên thưa hơn. Checkpoint đã huấn luyện lưu trạng thái thưa của SpongeTorch (các tensor _orig/_mask), vốn cần thiết cho bước xuất model sau đó. File cấu hình được sao chép vào thư mục lần chạy với tên amba_config.prototxt để bảo đảm khả năng tái lập.

Xác thực checkpoint đã nén#

Xác thực độ chính xác trước khi biên dịch bằng cùng cấu hình:

yolo val model=runs/detect/train/weights/best.pt data=coco8.yaml \
  amba_config=config.prototxt amba_chipset=CV72

Validator áp dụng lại spongetorch.prepare() khi cần và tắt việc hợp nhất Conv+BN để giữ nguyên cấu trúc nén. So sánh mAP với baseline chưa nén; nếu mức giảm độ chính xác quá lớn, hãy điều chỉnh cấu hình SpongeKit rồi huấn luyện lại.

Xuất sang ONNX#

Xuất checkpoint đã nén bằng chính amba_config như khi huấn luyện:

from ultralytics import YOLO

model = YOLO("runs/detect/train/weights/best.pt")
model.export(
    format="onnx",
    amba_config="config.prototxt",
    amba_chipset="CV72",
)

Exporter dựng lại model, áp dụng lại spongetorch.prepare() bằng cấu hình của bạn, nạp lại trọng số checkpoint thưa vào cấu trúc đã chuẩn bị và trace sang ONNX khi tắt hợp nhất Conv+BN — tạo ra đồ thị đúng định dạng mà công cụ biên dịch SDK yêu cầu.

Giữ lại metadata của model#

Quá trình xuất ONNX nhúng tác vụ của model, tên class, stride và kích thước đầu vào vào file ONNX, trong khi backend AmbaPB đọc thông tin này từ file sidecar metadata.yaml nằm cạnh model đã biên dịch. Nếu công cụ biên dịch SDK không tạo file sidecar này, hãy trích xuất metadata từ model ONNX trước khi biên dịch:

import onnx

from ultralytics.utils import YAML

model = onnx.load("model.onnx")
YAML.save("metadata.yaml", {item.key: item.value for item in model.metadata_props})

Giữ metadata.yaml trong cùng thư mục với file *.ambapb.ckpt.onnx hoặc *.ambapb.fastckpt.onnx đã biên dịch.

Biên dịch bằng SDK Tools#

Biên dịch model ONNX đã xuất cho chipset mục tiêu bằng các công cụ biên dịch SDK, theo hướng dẫn biên dịch của SDK. Các công cụ ánh xạ graph lên bộ xử lý AI CVflow — áp dụng PTQ, lập lịch và hoạch định bộ nhớ — rồi tạo checkpoint AmbaPB để xác thực trên host.

PTQ áp dụng lượng tử hóa INT8 bằng các ảnh hiệu chuẩn (chuẩn bị theo mô tả trong hướng dẫn biên dịch của SDK), còn các công cụ biên dịch cân bằng độ chính xác với độ trễ khi chạy: ánh xạ nhiều phép toán hơn sang INT8 giúp giảm độ trễ nhưng có thể làm giảm độ chính xác, trong khi giữ nhiều phép toán hơn ở FP16 sẽ bảo toàn độ chính xác với độ trễ cao hơn. Khi PTQ không thể đạt mục tiêu độ chính xác ở mức INT8 cần thiết cho ngân sách độ trễ, QAT với SpongeTorch là giải pháp được thiết kế để xử lý vấn đề này — phương pháp này huấn luyện model để chịu được lượng tử hóa INT8 mạnh hơn, khôi phục độ chính xác ở điểm vận hành có độ trễ thấp hơn. QAT hiện chưa có trong tích hợp này và được lên kế hoạch cho bản phát hành trong tương lai.

Chạy suy luận với model đã biên dịch#

Model AmbaPB đã biên dịch được tải trực tiếp qua API Ultralytics — AutoBackend phát hiện hậu tố .ambapb và chuyển hướng suy luận qua cvflowbackend, thực thi model theo đúng cách model sẽ chạy trên bộ xử lý AI:

from ultralytics import YOLO

model = YOLO("model.ambapb.ckpt.onnx")

# Inference
results = model("https://ultralytics.com/images/bus.jpg")

# Validation
metrics = model.val(data="coco8.yaml")

Đây là bước kiểm tra độ chính xác cuối cùng trước khi triển khai lên phần cứng, bao gồm mọi ảnh hưởng từ quá trình lượng tử hóa của trình biên dịch. Nếu có file metadata.yaml nằm cạnh model đã biên dịch, backend sẽ đọc tên lớp, stride và thông tin tác vụ từ file đó. Backend mặc định sử dụng chế độ suy luận CVflow acinf; đặt biến môi trường ULTRALYTICS_AMBAPB_DEBUG=1 để ghi nhật ký chi tiết đầu vào/đầu ra phục vụ gỡ lỗi.

Chuyển đổi sang định dạng nhị phân Cavalry#

Sau khi checkpoint AmbaPB vượt qua bước xác thực trên host, hãy dùng các công cụ biên dịch SDK để chuyển đổi checkpoint thành file nhị phân Cavalry cho thiết bị mục tiêu, theo hướng dẫn biên dịch của SDK. File nhị phân Cavalry là định dạng được thư viện runtime SDK thực thi trên bo mạch.

Triển khai trên bo mạch#

Tải file nhị phân Cavalry lên thiết bị Ambarella bằng runtime SDK Ambarella. Tiền xử lý và hậu xử lý phải khớp với cấu hình biên dịch của model phát hiện: đầu vào RGB đã letterbox với giá trị trong phạm vi 0–255 và giải mã đầu ra phát hiện YOLO tiêu chuẩn. Tham khảo tài liệu triển khai SDK để biết các API runtime.

Xuất không dùng SpongeTorch#

Nếu không cần tính năng tỉa trong quá trình huấn luyện của SpongeTorch, pipeline Ultralytics tiêu chuẩn cũng có thể tạo model để các công cụ SDK biên dịch:

yolo export model=yolo26n.pt format=onnx

Biên dịch ONNX thu được bằng các công cụ biên dịch SDK; các công cụ này tự thực hiện lượng tử hóa sau huấn luyện. Phương án này đánh đổi một phần hiệu năng khi chạy và độ chính xác sau lượng tử hóa để có quy trình đơn giản hơn, không phụ thuộc vào spongetorch trong quá trình huấn luyện.

Ứng dụng thực tế#

Các model Ultralytics YOLO trên SoC Ambarella CVflow cung cấp khả năng thị giác luôn hoạt động ở biên:

• Camera an ninh AI: phát hiện người và phương tiện theo thời gian thực trên camera IP 4K trong ngân sách công suất dưới 3 W.
• Drone và robot: phát hiện và theo dõi đối tượng trên thiết bị để phục vụ điều hướng, kiểm tra và giao hàng trên chip cấp CV5.
• Phân tích công nghiệp và bán lẻ: đếm người trên nhiều luồng, phát hiện PPE và giám sát kệ hàng trên các thiết bị edge.

Tóm tắt#

Hướng dẫn này trình bày quy trình hiện tại để triển khai model Ultralytics YOLO trên SoC Ambarella CVflow: huấn luyện có xét đến nén bằng SpongeTorch (amba_config/amba_chipset), xuất checkpoint đã nén sang ONNX, biên dịch ngoại tuyến thành checkpoint AmbaPB bằng công cụ SDK, xác thực trên host qua Ultralytics và chuyển đổi thành file nhị phân Cavalry để triển khai trên thiết bị bằng SDK Ambarella.

Để triển khai trên các mục tiêu AI biên khác, xem hướng dẫn liên quan về Hailo, Rockchip RKNN, Sony IMX500, Qualcomm QNN, DEEPX và Axelera. Để xem danh sách đầy đủ các định dạng xuất, hãy truy cập tài liệu chế độ Export và trang tích hợp.

Câu hỏi thường gặp#