如何使用 PyTorch 在 GPU 中运行特定代码

Question

我在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())

任何人都可以建议我如何解决这个问题吗？

Answer 1

你的类RoadShoulderWidth是一个nn。模块子类，允许您使用 .to(device(。这仅意味着所有其他 nn。模块对象或 nn。作为 RoadShoulderWidth 对象成员的参数将移动到设备。从您的示例中，没有，因此没有任何反应。

一般来说，PyTorch 不会将代码移动到 GPU，而是将数据。如果 pytorch 操作的所有数据都在 GPU 上(例如 a + b、a 和 b 在 GPU 上(，则该操作在 GPU 上执行。您可以使用 a.to(设备(移动数据，给定一个是割炬。张量对象。

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