You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

使用SSD与OpenCV进行目标检测时出现Runtime Error求助

解决PyTorch加载SSD模型时的CUDA/CPU不匹配RuntimeError

嘿,我来帮你搞定这个问题!

错误原因解释

这个报错的核心逻辑很直白:你用的预训练SSD模型文件ssd300_mAP_77.43_v2.pth是在带CUDA的GPU设备上训练并保存的,但你的当前运行环境没有可用的CUDA(也就是只能靠CPU跑代码)。PyTorch加载模型时默认会尝试把参数还原到保存时的设备(这里是CUDA),可你的机器没CUDA支持,自然就触发了这个RuntimeError。

修正后的完整代码

我只修改了模型加载的关键行,还顺手修正了一处拼写小错误,下面是修复后的完整脚本:

# Object Detection
# Importing Libraries
import torch
from torch.autograd import Variable
import cv2
from data import BaseTransform, VOC_CLASSES as labelmap
from ssd import build_ssd
import imageio

# Defining the function that does the Detection
def detect(frame, net, transform):
    height, width = frame.shape[:2]
    frame_t = transform(frame)[0]
    x = torch.from_numpy(frame_t).permute(2, 0, 1)
    x = Variable(x.unsqueeze(0))
    y = net(x)
    detections = y.data
    scale = torch.Tensor([width, height, width, height])
    for i in range(detections.size(1)):
        j = 0
        while detections[0, i, j, 0] >= 0.6:
            pt = (detections[0, i, j, 1:] * scale).numpy()
            cv2.rectangle(frame, (int(pt[0]), int(pt[1])), (int(pt[2]), int(pt[3])), (255, 0 , 0), 2)
            cv2.putText(frame, labelmap[i - 1], (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 2, cv2.LINE_AA)
            j += 1
    return frame

# Creating SSD Neural Networks
net = build_ssd('test')
# 关键修改:明确指定将模型加载到CPU设备
net.load_state_dict(torch.load('ssd300_mAP_77.43_v2.pth', map_location=torch.device('cpu')))

# Creating transformation
transform = BaseTransform(net.size, (104/256.0, 117/256.0, 123/256.0))

# Doing some Object Detection in the video
reader = imageio.get_reader('funny_dog.mp4')
fps = reader.get_meta_data()['fps']
writer = imageio.get_writer('output.mp4', fps=fps)
for i, frame in enumerate(reader):
    frame = detect(frame, net.eval(), transform)
    writer.append_data(frame)
    print(i)
writer.close()

核心修改说明

我把原来的模型加载代码:

net.load_state_dict(torch.load('ssd300_mAP_77.43_v2.pth'), map_location = lambda storage, loc: storage)

替换成了PyTorch官方推荐的写法:

net.load_state_dict(torch.load('ssd300_mAP_77.43_v2.pth', map_location=torch.device('cpu')))

这种写法更清晰地告诉PyTorch:不管模型原本是保存在什么设备上,都把它加载到当前的CPU设备上,完美适配你的CPU-only运行环境。

另外还修正了函数定义里的拼写错误fuction→function,不影响运行但更规范。现在运行这个脚本,应该就能正常执行视频目标检测了!

内容的提问来源于stack exchange,提问作者user12099710

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.05.08 19:27:40