如何使用PyTorch在GPU中运行垂直代码?

时间:2019-04-24 13:15:00

标签: multiprocessing gpu pytorch

我在python opencv中使用图像处理代码。由于该过程要花费大量时间来处理30张图像。我尝试使用多处理并行处理这些图像。多处理部分在CPU中运行良好,但是我想在GPU(cuda)中使用该多处理对象。

我使用torch.multiprocessing并行运行任务。因此,我在课堂上使用了torch.device('cuda')将整个设备运行到该垂直设备中。当我运行代码时,它使用“ cuda”显示设备,但不使用任何GPU处理。

import cv2
import numpy as np
import torch
import torch.nn as nn
from torch.multiprocessing import Process, Pool, Manager, set_start_method
import sys
import os

class RoadShoulderWidth(nn.Module):
    def __init__(self):  
        super(RoadShoulderWidth, self).__init__()
        pass

    // Want to run below method in parallel for 30 images.
    @staticmethod   
    def get_dim(image, road_shoulder_width_list):
        ..... code

    def get_road_shoulder_width(self, _root_dir, _img_path_list):

    manager = Manager()
    road_shoulder_width_list = manager.list()
    processes = []
    for img_path in img_path_list[:30]:
        img = cv2.imread(_root_dir + '/' + img_path)
        img = img[72 * 5:72 * 6, 0:1280]
        # Do work
        p = Process(target=self.get_dim,args=(img,road_shoulder_width_list))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
    return road_shoulder_width_list 

使用下面的代码集运行您的课程

if __name__ == '__main__':
    root_dir = '/home/nikhil_m/r'
    img_path_list = os.listdir(root_dir)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print('Using device:', device)
    dataloader_kwargs = {'pin_memory': True} 
    set_start_method('fork')   
    obj = RoadShoulderWidth().to(device)
    val = obj.get_road_shoulder_width(str(root_dir), img_path_list)
    print(val)
    print(torch.cuda.is_available())

有人可以建议我如何解决此问题吗?

1 个答案:

答案 0 :(得分:0)

您的类RoadShoulderWidth是nn.Module子类,可让您使用.to(device)。这仅意味着将作为RoadShoulderWidth对象成员的所有其他nn.Module对象或nn.Parameter移至设备。从您的示例来看,没有,所以什么也没发生。

通常,PyTorch不会将代码移至GPU,而是将数据移至GPU。如果pytorch操作的所有数据都在GPU上(例如a + b,a和b在GPU上),则该操作将在GPU上执行。如果a是torch.Tensor对象,则可以使用a.to(device)移动数据。

PyTorch只能在GPU上执行其自身的操作。它无法在GPU上执行OpenCV代码。