[FIXED] Why is the actual runtime still cpu after I set cuda for tensor

Issue

I want to gpu-accelerate a custom function, and below is the original function:

import numpy as np

def py_cpu_nms(dets, thresh):
    """Pure Python NMS baseline."""
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]

    return keep

It would use the CPU to compute, but that wasn't fast enough so I wanted to speed it up directly with PyTorch, which I converted into a Torch implementation:

import numpy as np
import torch

def py_cpu_nms(dets, thresh):
    """Pure Python NMS baseline."""
    dets = torch.from_numpy(dets)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    dets.cuda()
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = ((x2 - x1 + 1) * (y2 - y1 + 1)).cuda()
    order = torch.argsort(scores)
    order = torch.flip(order, dims=[0])
    keep = []
    while order.size()[0] > 0:
        i = order[0]
        keep.append(i)
        xx1 = torch.maximum(x1[i], x1[order[1:]]).cuda()
        yy1 = torch.maximum(y1[i], y1[order[1:]]).cuda()
        xx2 = torch.minimum(x2[i], x2[order[1:]]).cuda()
        yy2 = torch.minimum(y2[i], y2[order[1:]]).cuda()

        w = torch.maximum(torch.tensor(0.0), xx2 - xx1 + 1).cuda()
        h = torch.maximum(torch.tensor(0.0), yy2 - yy1 + 1).cuda()
        inter = (w * h).cuda()
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = torch.where(ovr <= thresh)[0].cuda()
        order = order[inds + 1].cuda()

    return keep

But in fact, all the calculations still use the CPU, does anyone know why?

Solution

cuda() is no in-place. Change your line 3 of the function to dets = dets.cuda().

Answered By - joe32140