如何使用PyTorch中Pythonic风格的YOLOv5完成训练与推理

你提到的Pythonic调用方式本质是通过torch.hub加载YOLOv5模型的封装接口，训练及预测结果提取的代码示例如下：

训练代码示例 #

import torch

# 加载预训练YOLOv5s模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=True)

# 启动训练，按需调整参数
model.train(
    data='你的数据集yaml配置文件路径',
    epochs=50,
    imgsz=640,
    batch_size=16,
    name='训练任务自定义名称'
)

预测结果提取代码示例 #

# 加载训练完成的权重文件
model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/训练任务自定义名称/weights/best.pt', autoshape=True)

# 单张图片推理
results = model('test.jpg')
# 提取结果：格式为[xmin, ymin, xmax, ymax, 置信度, 类别ID]
bbox_info = results.xyxy[0].cpu().numpy()

不会丢失任何核心功能。你提到的官方文档里的非Pythonic方式，本质是命令行调用train.py/detect.py脚本的用法，Pythonic接口是对这些脚本核心逻辑的原生封装，训练、推理、评估、模型导出所有官方支持的功能，都可以通过接口参数配置实现，仅有的区别是不需要通过命令行传参，直接在Python代码中传入对应参数即可。

不需要修改Roboflow导出的数据格式。Roboflow导出的YOLOv5格式数据集（包含images文件夹、labels文件夹、对应数据集yaml配置文件）完全适配Pythonic接口的训练要求，直接把yaml文件路径传入model.train()的data参数即可正常使用。

单张图片推理确实可以直接写prediction = model('my_image.jpg')，除了本地图片路径，你还可以直接传入PIL图片对象、OpenCV读取的numpy数组、PyTorch张量格式的图片数据。
多张图片预测直接将多个图片路径/图片对象放在列表中传入即可，示例如下：

# 多张图片推理示例
img_list = ['img1.jpg', 'img2.jpg', 'img3.jpg']
results = model(img_list)

# 遍历提取每张图片的预测结果
for idx in range(len(img_list)):
    single_img_result = results.xyxy[idx].cpu().numpy()
    # 后续自定义处理逻辑

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