无Ultralytics环境下YOLOv8模型导出至TF/PyTorch方案咨询

一、导出为纯PyTorch可推理格式 #

你可以直接导出为TorchScript格式，这种格式的模型无需依赖Ultralytics包，仅用PyTorch就能完成推理：

1. 在有Ultralytics的环境中加载并导出模型：

from ultralytics import YOLO
import torch

model = YOLO("{HOME}/runs/detect/train/weights/best.pt")
model.fuse()  # 融合层提升推理效率
ts_model = torch.jit.trace(model.model, torch.randn(1, 3, 1024, 1024))
ts_model.save("yolov8_best_ts.pt")

2. 在仅装PyTorch的目标服务器上推理：

import torch
import numpy as np
from PIL import Image

# 加载模型
model = torch.jit.load("yolov8_best_ts.pt")
model.eval()

# 图片预处理
img = Image.open("test.jpg").resize((1024, 1024))
img_tensor = torch.tensor(np.array(img)).permute(2, 0, 1).unsqueeze(0).float() / 255.0

# 推理
with torch.no_grad():
    outputs = model(img_tensor)

# 解析输出（格式为[batch, num_boxes, 5+num_classes]，前5项为x,y,w,h,置信度）
boxes = outputs[0][..., :4]
scores = outputs[0][..., 4]
class_ids = outputs[0][..., 5:].argmax(dim=-1)

二、导出为TensorFlow/Keras格式 #

由于Ultralytics 8.0.20的官方TF导出功能未完成，可以通过ONNX中间格式转换：

1. 在有Ultralytics的环境中导出ONNX模型：

model.export(format="onnx", imgsz=1024)
# 导出后生成best.onnx文件

2. 安装转换工具并将ONNX转成TensorFlow SavedModel：

pip install onnx-tf

import onnx
from onnx_tf.backend import prepare

onnx_model = onnx.load("best.onnx")
tf_rep = prepare(onnx_model)
tf_rep.export_graph("yolov8_tf_savedmodel")

3. 在仅装TensorFlow的目标服务器上推理：

import tensorflow as tf
from PIL import Image
import numpy as np

# 加载模型
model = tf.saved_model.load("yolov8_tf_savedmodel")
infer = model.signatures["serving_default"]

# 图片预处理
img = Image.open("test.jpg").resize((1024, 1024))
img_array = np.array(img).astype(np.float32) / 255.0
img_tensor = tf.convert_to_tensor(img_array[np.newaxis, ...])
# 若输入格式要求为NCHW，执行：img_tensor = tf.transpose(img_tensor, [0, 3, 1, 2])

# 推理
outputs = infer(img_tensor)
output_tensor = outputs["output0"]
boxes = output_tensor[..., :4]
scores = output_tensor[..., 4]
class_ids = tf.argmax(output_tensor[..., 5:], axis=-1)

内容的提问来源于stack exchange，提问作者ilia0017