python - 如何将基于 pytorch cpu 的转换转换为基于 cuda 的转换？

Question

代码的原始问题已发布 here .

我正在使用this线分割项目的存储库，我开发了此代码来获取输入(无论是图像还是视频)并在其上绘制道路线并将其提供给输出:

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] > 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] > 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()

Resize、ToTensor、Normalize、Compose 的代码在这里:

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

该代码工作正常，但我发现它对于在实时应用程序中进行测试来说太慢了。我添加了一些时间测量来查看是否可以找出瓶颈，这是一个循环的输出:

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

这些行恰好是计算成本最高的行:

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]

现在我遇到了这个问题，即如何修改代码以提高计算速度。

最初我想修改代码以允许cuda计算。我问主要作者怎么修改here中cuda版本的代码他指出了这些行:

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

不幸的是我对 pytorch 的经验不多，所以我现在寻求帮助。

我希望我分享的信息足以满足读者的需求。任何帮助将不胜感激

谢谢

Answer 1

设置设备:

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

他将数据放在设备上的意思是:

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)