File size: 2,987 Bytes
157c353 | 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 | #!/usr/bin/env python3
import torch
from ultralytics import YOLO
from ultralytics.nn.modules.head import Pose
import os
import shutil
def npu_pose_forward(self, x):
"""
YOLOv8 Pose Head Modified for NPU.
YOLOv8 Pose structure (NOT end2end):
- cv2: box regression layers (DFL, out channels = 4 * reg_max = 64)
- cv3: classification layers (out channels = nc, nc=1 for pose)
- cv4: keypoint output layers (out channels = nk = 17 * 3 = 51)
Output:
List of Tensors (9 items for 3 scales), Layout: NHWC.
For each scale:
- Box_Raw (B, H, W, 4*reg_max), <-- raw DFL box logits (decode in post-process)
- Cls_Raw (B, H, W, nc), <-- class scores (raw logits), nc=1 for pose
- Kpt_Raw (B, H, W, nk), <-- raw keypoints (nk = 17*3 = 51)
"""
if not isinstance(x, (list, tuple)):
x = [x]
res = []
box_layers = self.cv2
cls_layers = self.cv3
kpt_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. Keypoints branch (raw) - NHWC
kpts = kpt_layers[i](x[i]).permute(0, 2, 3, 1)
res.append(bboxes)
res.append(scores)
res.append(kpts)
return res
def batch_export_yolov8_pose():
variants = ['n', 's', 'm', 'l', 'x']
imgsz = 640
# Execute Monkey Patch
Pose.forward = npu_pose_forward
print("Monkey patch applied for Pose: Output Layout forced to NHWC (Box, Cls, Kpt for each scale).")
for v in variants:
model_name = f"yolov8{v}-pose"
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
Pose.forward = npu_pose_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, Pose):
head.forward = lambda x: npu_pose_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_pose()
|