Fix and update YoloV6 Nano proof point

demo_helpers/demo_helpers/model_download.py

@@ -1,19 +1,23 @@
 import os
+import zipfile
 
 from datasets.utils.file_utils import cached_path
 from groqflow.common.build import DEFAULT_CACHE_DIR
 
 
-YOLOV6N_ONNX = "yolov6n_onnx"
+YOLOV6N_MODEL = "yolov6n_model"
+YOLOV6N_SOURCE = "yolov6n_source"
 
 
 DATA_URLS = {
-    YOLOV6N_ONNX: "https://github.com/meituan/YOLOv6/releases/download/0.1.0/yolov6n.onnx",
+    YOLOV6N_MODEL: "https://github.com/meituan/YOLOv6/releases/download/0.4.0/yolov6n.pt",
+    YOLOV6N_SOURCE: "https://github.com/meituan/YOLOv6/archive/refs/tags/0.4.0.zip",
 }
 
 
 DST_PATHS = {
-    YOLOV6N_ONNX: "onnx_models/yolov6n.onnx",
+    YOLOV6N_MODEL: "pytorch_models/yolov6_nano/yolov6n.pt",
+    YOLOV6N_SOURCE: "pytorch_models/yolov6_nano/YOLOv6",
 }
 
 
@@ -27,3 +31,18 @@ def download_model(model):
     download_path = cached_path(url)
     os.symlink(download_path, dst_path)
     return dst_path
+
+
+def download_source(source):
+    dst_path = os.path.join(DEFAULT_CACHE_DIR, DST_PATHS[source])
+    if os.path.exists(dst_path):
+        return dst_path
+
+    os.makedirs(os.path.dirname(dst_path), exist_ok=True)
+    url = DATA_URLS[source]
+    download_path = cached_path(url)
+    with zipfile.ZipFile(download_path, "r") as zip_ref:
+        extracted_dir = os.path.dirname(dst_path)
+        zip_ref.extractall(extracted_dir)
+        os.rename(os.path.join(extracted_dir, zip_ref.infolist()[0].filename), dst_path)
+    return dst_path

demo_helpers/demo_helpers/models.py

@@ -1,8 +1,18 @@
 import os
+import subprocess
+import sys
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from demo_helpers.model_download import (
+    YOLOV6N_MODEL,
+    YOLOV6N_SOURCE,
+    download_model,
+    download_source,
+)
+
 
 class M5(nn.Module):
     def __init__(self, n_input=1, n_output=35, stride=16, n_channel=32):
@@ -130,6 +140,33 @@ def forward(self, input):
         return self.logsoftmax(output)
 
 
+def get_yolov6n_model():
+    weights = download_model(YOLOV6N_MODEL)
+    source = download_source(YOLOV6N_SOURCE)
+    export_script = os.path.join(source, "deploy/ONNX/export_onnx.py")
+
+    cmd = [
+        sys.executable,
+        export_script,
+        "--weights",
+        weights,
+        "--img",
+        "640",
+        "--batch",
+        "1",
+        "--simplify",
+    ]
+    p = subprocess.Popen(
+        cmd, cwd=source, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    )
+    p.communicate()
+    if p.returncode != 0:
+        raise RuntimeError("Unable to get ONNX model")
+
+    onnx_file = weights.replace(".pt", ".onnx")
+    return onnx_file
+
+
 def load_pretrained(model_name):
     """Loads a pre-trained model
 

docs/known_issues.md → docs/release_notes.md

@@ -1,4 +1,22 @@
-# GroqFlow Known Issues
+# Release Notes
+
+## v4.3.0
+
+### Changes
+
+* Support for SDK 0.11.
+* Add beta support for groq-torch-importer front-end support.
+* Clean up package dependencies.
+* Various bug fixes.
+
+### Known Issues
+
+* Yolo V6 proof point downloads the pytorch weights and invokes the export script to get the ONNX file.
+* Pip install of GroqFlow may complain about incompatible protobuf version.
+
+## v4.2.1
+
+### Known Issues
 
 * Runtime errors due to mismatches in tensor sizes may occur even though GroqFlow checks the data shape. (G14148)
 * Whacky terminal line wrapping when printing groqit error messages. (G13235)

proof_points/computer_vision/yolo/yolov6_nano.py

@@ -4,33 +4,29 @@
 the COCO dataset (https://cocodataset.org/) on CPU and
 GroqChip™ processor using the GroqFlow toolchain.
 """
+import torch
 
 from groqflow import groqit
 from demo_helpers.args import parse_args
 from demo_helpers.compute_performance import compute_performance
-from demo_helpers.model_download import YOLOV6N_ONNX, download_model
+from demo_helpers.models import get_yolov6n_model
 from demo_helpers.misc import check_deps
 
-import torch
-
-
-def get_onnx_model():
-    return download_model(YOLOV6N_ONNX)
-
 
 def evaluate_yolov6n(rebuild_policy=None, should_execute=True):
     check_deps(__file__)
-    pytorch_model = get_onnx_model()
-    dummy_inputs = {"image_arrays": torch.ones([1, 3, 640, 640])}
+    model = get_yolov6n_model()
+    dummy_inputs = {"images": torch.ones([1, 3, 640, 640])}
 
     # Get Groq Model using groqit
     groq_model = groqit(
-        pytorch_model,
+        model,
         dummy_inputs,
         rebuild=rebuild_policy,
+        compiler_flags=["--effort=high"],
     )
     if should_execute:
-        compute_performance(groq_model, pytorch_model, "coco", task="coco_map")
+        compute_performance(groq_model, model, "coco", task="coco_map")
 
     print(f"Proof point {__file__} finished!")