ayooshkathuria / YOLO_v3_tutorial_from_scratch

Accompanying code for Paperspace tutorial series "How to Implement YOLO v3 Object Detector from Scratch"
https://blog.paperspace.com/how-to-implement-a-yolo-object-detector-in-pytorch/
2.32k stars 724 forks source link

RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _thnn_conv2d_forward #68

Open grok0n opened 3 years ago

grok0n commented 3 years ago

I am trying to run the forward pass on my GPU but it is not working. I wrote a detector class and when I try to run the forward pass of the model on GPU I get the above error.

import cv2
from darknet import Darknet
from util import *
from torch.autograd import Variable
import torch

class ObjectDetector:
    def __init__(self, confidence=0.5, nms_thresh=0.4, reso=416):

        #Load Darknet model for object detection
        print("Loading network...")
        self.model = Darknet("cfg/yolov3.cfg")
        self.model.load_weights("yolov3.weights")
        print("Network successfully loaded")

        self.model.net_info["height"] = reso
        self.inp_dim = int(self.model.net_info["height"])
        if torch.cuda.is_available():
            self.model.cuda()

        self.model.eval()

        #Set detector parameters
        self.confidence = confidence
        self.nms_thresh = nms_thresh
        self.reso = reso

        #Load class file
        self.num_classes = 80
        self.classes = load_classes("data/coco.names")

    def detect(self, img):
        img_ = prep_image(img, self.inp_dim)

        if torch.cuda.is_available():
            img_.cuda()

        pred = self.model(img_, torch.cuda.is_available())
        res = write_results(pred, self.confidence, self.num_classes)[:,[1,2,3,4,7]] # only take the 4 corner coordinate points and class index

        # Create a python list for all the objects with the format [x1,y1,x2,y2,class_name]
        objs = []
        for obj in res:
            _ = obj[:4].tolist()
            _.append(self.classes[int(obj[4])])
            objs.append(_)

        return objs

if __name__=="__main__":
    detector = ObjectDetector()
    img = cv2.imread(input("Input image path: "))
    detector.detect(img)

I works well on the CPU though.