How can I convert pytorch cpu-based transformation to cuda-based?
The original issue for the code is availablehere.

I am using this repository for a line segmentation project and I developed this code to get an input (whether image or video) and draw road lines on it and give it in output:

import argparse
import sys
from time import time, clock
from os.path import splitext, basename, exists

from model import SCNN
from utils.check_extension import is_video, is_image
from utils.transforms import *
# I will put all the necessary code for utils.transforms after this

# ------------------------------------------------  SCNN parameters
time1 = time()
net = SCNN(input_size=(800, 288), pretrained=False)
mean = (0.3598, 0.3653, 0.3662)  # CULane mean, std
std = (0.2573, 0.2663, 0.2756)
transform_img = Resize((800, 288))
transform_to_net = Compose(ToTensor(), Normalize(mean=mean, std=std))


# ------------------------------------------------  Arguments


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--weights', type=str,
                        default='models/vgg_SCNN_DULR_w9.pth',
                        help='path to vgg models')

    parser.add_argument('--input', type=str, default='demo/line_3.mp4',
                        help='path to image file')

    parser.add_argument('--output', type=str, default='public/',
                        help='path to the output directory')
    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    filename, extension = splitext(basename(args.input))

    print("Loading file [{}] ....".format(filename))
    if not exists(args.input):
        print("file [{}] is not recognized".format(args.input))
        sys.exit()


    if is_video(extension):
        video_capture = cv2.VideoCapture()
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        output = args.output + filename + '.avi'

        if video_capture.open(args.input):
            property_id = int(cv2.CAP_PROP_FRAME_COUNT)
            total_frames = int(cv2.VideoCapture.get(video_capture, property_id))
            frame_no = 1
            width, height = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)), \
                            int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
            fps = video_capture.get(cv2.CAP_PROP_FPS)
            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            save_dict = torch.load(args.weights, map_location=device)
            net.load_state_dict(save_dict['net'])
            net.eval()

            # can't write out mp4, so try to write into an AVI file
            video_writer = cv2.VideoWriter(output, fourcc, fps, (width, height))
            while video_capture.isOpened():
                start = time()
                ret, frame = video_capture.read()
                if not ret:
                    break
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frame = transform_img({'img': frame})['img']
                x = transform_to_net({'img': frame})['img']
                x.unsqueeze_(0)

                stop1 = time()
                print('stop1: ', stop1 - start)

                seg_pred, exist_pred = net(x)[:2]
                seg_pred = seg_pred.detach().cpu().numpy()
                exist_pred = exist_pred.detach().cpu().numpy()
                seg_pred = seg_pred[0]

                stop2 = time()
                print('stop2: ', stop2 - stop1)

                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                lane_img = np.zeros_like(frame)
                color = np.array([[255, 125, 0], [0, 255, 0], [0, 0, 255], [0, 255, 255]], dtype='uint8')
                coord_mask = np.argmax(seg_pred, axis=0)
                for i in range(0, 4):
                    if exist_pred[0, i] &gt; 0.5:
                        lane_img[coord_mask == (i + 1)] = color[i]
                img = cv2.addWeighted(src1=lane_img, alpha=0.8, src2=frame, beta=1., gamma=0.)
                img = cv2.resize(img, (width, height))

                stop3 = time()
                print('stop3: ', stop3 - stop2)

                # if frame_no % 20 == 0:
                #     print('# {}/{} frames processed!'.format(frame_no, total_frames))
                frame_no += 1
                video_writer.write(img)
                end = time()
                print('Whole loop: {} seconds'.format(end - start))
                print('------------')
                print('------------')

            print('# All frames processed ')

            video_capture.release()
            video_writer.release()


    elif is_image(extension):
        img = cv2.imread(args.input)
        height, width, _ = img.shape
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = transform_img({'img': img})['img']
        x = transform_to_net({'img': img})['img']
        x.unsqueeze_(0)


        save_dict = torch.load(args.weights, map_location='cpu')
        net.load_state_dict(save_dict['net'])
        net.eval()

        seg_pred, exist_pred = net(x)[:2]
        seg_pred = seg_pred.detach().cpu().numpy()
        exist_pred = exist_pred.detach().cpu().numpy()
        seg_pred = seg_pred[0]


        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        lane_img = np.zeros_like(img)
        color = np.array([[255, 125, 0], [0, 255, 0], [0, 0, 255], [0, 255, 255]], dtype='uint8')
        coord_mask = np.argmax(seg_pred, axis=0)
        for i in range(0, 4):
            if exist_pred[0, i] &gt; 0.5:
                lane_img[coord_mask == (i + 1)] = color[i]
        img = cv2.addWeighted(src1=lane_img, alpha=0.8, src2=img, beta=1., gamma=0.)
        img = cv2.resize(img, (width, height))
        output = args.output + filename + '.jpg'

        cv2.imwrite(output, img)

    else:
        print("file format [{}] is not supported".format(args.input))
        sys.exit()


if __name__ == '__main__':
    main()


The code which belong to Resize, ToTensor, Normalize, Compose are here:

class Compose(CustomTransform):
    """
    All transform in Compose should be able to accept two non None variable, img and boxes
    """
    def __init__(self, *transforms):
        self.transforms = [*transforms]

    def __call__(self, sample):
        for t in self.transforms:
            sample = t(sample)
        return sample

    def __iter__(self):
        return iter(self.transforms)

    def modules(self):
        yield self
        for t in self.transforms:
            if isinstance(t, Compose):
                for _t in t.modules():
                    yield _t
            else:
                yield t


class Normalize(CustomTransform):
    def __init__(self, mean, std):
        self.transform = Normalize_th(mean, std)

    def __call__(self, sample):
        img = sample.get('img')

        img = self.transform(img)

        _sample = sample.copy()
        _sample['img'] = img
        return _sample


class ToTensor(CustomTransform):
    def __init__(self, dtype=torch.float):
        self.dtype=dtype

    def __call__(self, sample):
        img = sample.get('img')
        segLabel = sample.get('segLabel', None)
        exist = sample.get('exist', None)

        img = img.transpose(2, 0, 1)
        img = torch.from_numpy(img).type(self.dtype) / 255.
        if segLabel is not None:
            segLabel = torch.from_numpy(segLabel).type(torch.long)
        if exist is not None:
            exist = torch.from_numpy(exist).type(torch.float32)  # BCEloss requires float tensor

        _sample = sample.copy()
        _sample['img'] = img
        _sample['segLabel'] = segLabel
        _sample['exist'] = exist
        return _sample


class Resize(CustomTransform):
    def __init__(self, size):
        if isinstance(size, int):
            size = (size, size)
        self.size = size  #(W, H)

    def __call__(self, sample):
        img = sample.get('img')
        segLabel = sample.get('segLabel', None)

        img = cv2.resize(img, self.size, interpolation=cv2.INTER_CUBIC)
        if segLabel is not None:
            segLabel = cv2.resize(segLabel, self.size, interpolation=cv2.INTER_NEAREST)

        _sample = sample.copy()
        _sample['img'] = img
        _sample['segLabel'] = segLabel
        return _sample

    def reset_size(self, size):
        if isinstance(size, int):
            size = (size, size)
        self.size = size


The code works fine but I found out that its too slow for testing in real-time application. I added some time measurement to see if I can find out the bottlenecks and this is the output for one loop:

------------
stop1:  0.002989053726196289
stop2:  1.4032211303710938
stop3:  0.004946708679199219
Whole loop: 1.41636061668396 seconds


These lines happened to be the most computationally expensive lines:

seg_pred, exist_pred = net(x)[:2]
seg_pred = seg_pred.detach().cpu().numpy()
exist_pred = exist_pred.detach().cpu().numpy()
seg_pred = seg_pred[0]


Now I am stuck with this issue that how I can modify the code to improve the computation speed. 

Initially I thought of modifying the code to allow cuda computation. I asked the main author how I can modify the code for cuda version in here and he pointed out to these lines:

frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame = transform_img({'img': frame})['img']
x = transform_to_net({'img': frame})['img']
x.unsqueeze_(0)


Unfortunately my experience with pytorch is not much, so I am asking for help now. 

I hope the information I shared suffices for the readers. Any help would be appreciated

Thanks

Set device:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


what he means his putting the data on device:

frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame = transform_img({'img': frame})['img']
x = transform_to_net({'img': frame})['img']
x.unsqueeze_(0).to(device)

