Apple M1 Pro运行YOLO-NAS推理视频遇topk RuntimeError求助

问题场景 #

在Apple M1 Pro设备上测试SuperGradients的YOLO-NAS模型，执行以下视频推理代码：

from super_gradients.training import models

model = models.get("yolo_nas_s_int8", pretrained_weights="coco")
model.to("mps")
model.predict("video_path") # 触发错误

运行后抛出错误：

RuntimeError: Currently topk on mps works only for k<=16

原因分析 #

• SuperGradients官方文档未明确标注支持MPS设备，model.to("mps")仅完成设备迁移，但模型内部推理逻辑中使用的topk操作超出了当前PyTorch MPS后端的限制（仅支持k≤16）
• 该问题是PyTorch MPS后端的已知限制，即使使用最新的PyTorch nightly版本（2.1.0.dev20230716）仍存在此问题
• Ultralytics YOLOv8可正常运行MPS，是因为其推理逻辑中对topk的使用符合MPS后端的限制

解决方案 #

• 临时替代方案：切换到CPU或CUDA设备运行，修改代码为model.to("cpu")或model.to("cuda")（若有外接GPU）
• 等待官方修复：待PyTorch MPS后端支持更大k值的topk操作后，再尝试使用MPS运行
• 自定义修改（需熟悉模型代码）：手动调整YOLO-NAS推理阶段topk的k值至16以内，注意此操作可能降低检测结果的召回率

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