Add Yolo26 to matrix of tested models on RISC-V

.github/workflows/riscv64.yml

@@ -28,21 +28,22 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        include:
-          - { model: add,        xnnpack: false, quantize: false }
-          - { model: add,        xnnpack: true,  quantize: false }
-          - { model: mv2,        xnnpack: false, quantize: false }
-          - { model: mv2,        xnnpack: true,  quantize: false }
-          - { model: mv2,        xnnpack: true,  quantize: true }
-          - { model: mobilebert, xnnpack: false, quantize: false }
-          - { model: mobilebert, xnnpack: true,  quantize: false }
-          - { model: mobilebert, xnnpack: true,  quantize: true }
-          - { model: llama2,     xnnpack: false, quantize: false }
-          - { model: llama2,     xnnpack: true,  quantize: false }
-          - { model: llama2,     xnnpack: true,  quantize: true }
-          - { model: resnet18,   xnnpack: false, quantize: false }
-          - { model: resnet18,   xnnpack: true,  quantize: false }
-          - { model: resnet18,   xnnpack: true,  quantize: true }
+        model:
+          - add
+          - mv2
+          - mobilebert
+          - llama2
+          - resnet18
+          - yolo26
+        xnnpack: [true, false]
+        quantize: [true, false]
+        exclude:
+          # We only enable quantization with XNNPACK
+          - xnnpack: false
+            quantize: true
+          # We don't test quantization for Yolo26
+          - model: yolo26
+            quantize: true
     permissions:
       id-token: write
       contents: read

examples/riscv/aot_riscv.py

@@ -114,12 +114,45 @@ def build_resnet18():
     return model, example_inputs, test_inputs, False
 
 
+def build_yolo26():
+    # Mirrors examples/models/yolo26/export_and_validate.py: predict() once
+    # to materialise the predictor state Ultralytics expects pre-export.
+    import numpy as np
+    from ultralytics import YOLO
+
+    input_h, input_w = 320, 320
+    yolo = YOLO("yolo26n")
+    yolo.predict(
+        np.ones((input_h, input_w, 3)),
+        imgsz=(input_h, input_w),
+        device="cpu",
+    )
+
+    class Wrapper(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = yolo.model.to(torch.device("cpu")).eval()
+
+        def forward(self, x):
+            # yolo.model emits (predictions, feature_maps) in eval; keep the
+            # predictions tensor so BundledIO sees a single tensor output.
+            out = self.model(x)
+            return out[0] if isinstance(out, (tuple, list)) else out
+
+    model = Wrapper().eval()
+    torch.manual_seed(0)
+    example_inputs = (torch.randn(1, 3, input_h, input_w),)
+    test_inputs = [example_inputs]
+    return model, example_inputs, test_inputs, False
+
+
 MODELS = {
     "add": build_add,
     "mv2": build_mv2,
     "mobilebert": build_mobilebert,
     "llama2": build_llama2,
     "resnet18": build_resnet18,
+    "yolo26": build_yolo26,
 }
 
 

examples/riscv/requirements.txt

@@ -1,2 +1,3 @@
 torchvision
 transformers
+ultralytics

examples/riscv/setup.sh

@@ -33,7 +33,10 @@ ${SUDO} apt-get install -y --no-install-recommends \
     cmake \
     file \
     ca-certificates \
-    qemu-user-static
+    qemu-user-static \
+    libglib2.0-0t64 \
+    libxcb1 \
+    libgl1
 
 if [[ -n "${GCC_VERSION+x}" ]]; then
     ${SUDO} update-alternatives --install /usr/bin/riscv64-linux-gnu-gcc riscv64-linux-gnu-gcc /usr/bin/riscv64-linux-gnu-gcc${GCC_VERSION:+-${GCC_VERSION}} 100