File size: 3,549 Bytes
f3370c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | #!/usr/bin/env python3
import torch
from ultralytics import YOLO
from ultralytics.nn.modules.head import Segment
import os
import shutil
def npu_segment_forward(self, x):
"""
YOLOv8 Segment Head Modified for NPU.
YOLOv8 Segment structure (NOT end2end):
- cv2: box regression layers (DFL, out channels = 4 * reg_max = 64)
- cv3: classification layers (out channels = nc = 80)
- cv4: mask coefficient layers (out channels = nm = 32)
- proto: Proto module for generating mask prototypes
Output:
List of Tensors (10 items: 3 scales * 3 outputs + 1 proto), NHWC for detection, NCHW for proto
[
Scale1_Box_Raw (B, 80, 80, 4*reg_max), <-- raw DFL box logits
Scale1_Cls_Raw (B, 80, 80, nc), <-- class scores (raw logits), nc=80
Scale1_Mask_Raw (B, 80, 80, nm), <-- mask coefficients, nm=32
Scale2_Box_Raw (B, 40, 40, 4*reg_max),
Scale2_Cls_Raw (B, 40, 40, nc),
Scale2_Mask_Raw (B, 40, 40, nm),
Scale3_Box_Raw (B, 20, 20, 4*reg_max),
Scale3_Cls_Raw (B, 20, 20, nc),
Scale3_Mask_Raw (B, 20, 20, nm),
Proto (B, nm, H, W), <-- mask prototypes (160x160 for 640 input)
]
"""
if not isinstance(x, (list, tuple)):
x = [x]
res = []
box_layers = self.cv2
cls_layers = self.cv3
mask_layers = self.cv4
for i in range(self.nl):
# 1. Box branch (raw DFL logits) - NHWC
bboxes = box_layers[i](x[i]).permute(0, 2, 3, 1)
# 2. Cls branch (raw logits) - NHWC
scores = cls_layers[i](x[i]).permute(0, 2, 3, 1)
# 3. Mask coefficients branch - NHWC
masks = mask_layers[i](x[i]).permute(0, 2, 3, 1)
res.append(bboxes)
res.append(scores)
res.append(masks)
# 4. Proto output - NCHW (keep original format for mask processing)
proto = self.proto(x[0])
res.append(proto)
return res
def batch_export_yolov8_seg():
variants = ['n', 's', 'm', 'l', 'x']
imgsz = 640
# Execute Monkey Patch
Segment.forward = npu_segment_forward
print("Monkey patch applied for Segment: Output Layout forced to NHWC + Proto.")
for v in variants:
model_name = f"yolov8{v}-seg"
pt_path = f"{model_name}.pt"
onnx_final_name = f"{model_name}_640x640.onnx"
print(f"\n--- Processing {model_name} ---")
try:
# Load model
model = YOLO(pt_path)
# Reapply monkey patch
Segment.forward = npu_segment_forward
# Ensure the model's head also uses the new forward
if hasattr(model.model, 'model') and len(model.model.model) > 0:
head = model.model.model[-1]
if isinstance(head, Segment):
head.forward = lambda x: npu_segment_forward(head, x)
# Execute export
exported_path = model.export(
format="onnx",
imgsz=imgsz,
dynamic=False,
opset=11,
simplify=True,
nms=False
)
# Move and rename
if exported_path:
shutil.move(exported_path, onnx_final_name)
print(f"Success: {onnx_final_name}")
except Exception as e:
print(f"Failed to export {model_name}: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
batch_export_yolov8_seg()
|