YOLOv8-Pose / export_onnx.py

First release

157c353 about 2 months ago

2.99 kB

#!/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()

	#!/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()