pycuda pytorch交互——使用pycuda处理pytorch tensor

pycuda pytorch交互——使用pycuda处理pytorch tensor很多时候,我们希望对tensor有一些自定义的操作,一种实现方式就是使用pycuda。本文以实现两个tensor的加法为例,讲解如何实现pycuda与pytorch交互。1.pycuda的使用方式首先看下pycuda文档对pycuda的定义:PyCUDAgivesyoueasy,PythonicaccesstoNvidia’sCUDAparallelcomputationAPI.即pycuda是在python中调用CUDA进行并行计算的接口。使用这个接口的流程为:1

pycuda

很多时候,我们希望对tensor有一些自定义的操作,一种实现方式就是使用pycuda。本文以实现两个tensor的加法为例,讲解如何实现pycuda与pytorch交互。

1. pycuda的使用方式

首先看下pycuda文档对pycuda的定义:

PyCUDA gives you easy, Pythonic access to Nvidia’s CUDA parallel computation API.

即pycuda是在python中调用CUDA进行并行计算的接口。
使用这个接口的流程为:

1.初始化PyCUDA模块。
2.导入一个用于编译的类。
3.定义一个存储着CUDA代码的字符串。
4.构造一个类对象(编译CUDA代码,使之成为可以被执行的GPU程序)。
5.从类对象中获得一个函数的执行入口。
6.通过这个入口,执行GPU程序。

原理和CUDA c++一致,不太了解的同学可以先去看一下。不同的是kenel函数要定义在一个字符串中,通过构造类对象编译CUDA代码,然后统过类对象中的函数接口进行调用。
下面举一个简单的例子,实现了两个二维数组的相加:

import numpy as np
import pycuda.autoinit #以自动的方式对pycuda进行初始
from pycuda.compiler import SourceModule # 编译kernel函数的类
import pycuda.gpuarray as gpuarray
# 通过字符串定义kernel函数
kernel_code = r""" void __global__ add(const float *x, const float *y, float *z) { const int n = threadIdx.y*blockDim.x+threadIdx.x; z[n] = x[n] + y[n]; } """
# 编译kernel函数
mod = SourceModule(kernel_code)
# 获取函数接口
add= mod.get_function("add")
# 定义三个二维数组并转换成pycuda中的gpuarray
A = np.zeros((10, 10))
B = np.ones((10, 10))
C = np.zeros((10, 10))
A_GPU = gpuarray.to_gpu(A.astype(np.float32))
B_GPU = gpuarray.to_gpu(B.astype(np.float32))
C_GPU = gpuarray.to_gpu(B.astype(np.float32))
# 执行kernel函数,和cuda c++一样定义griddim和blockdim;
add(A_GPU, B_GPU, C_GPU, grid=(1,1,1), block=(10,10,1))
C = C_GPU.get()
print(C)

输出:

[[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]]

编译kernel代码的时间比较耗时,编译一次之后后面就快了。

2. pycuda与pytorch交互处理tensor

要想完成两者的交互,需要定义以下这个类,并利用tensor初始化一个对象。

class Holder(pycuda.driver.PointerHolderBase):
    def __init__(self, t):
        super(Holder, self).__init__()
        self.t = t
        self.gpudata = t.data_ptr()
    def get_pointer():
        return self.t.data_ptr()

还是上面的二维矩阵相加的例子,但A,B,C换成pytorch的tensor:

import torch
import pycuda.autoinit
from pycuda.compiler import SourceModule
x = torch.cuda.FloatTensor(8)
class Holder(pycuda.driver.PointerHolderBase):
    def __init__(self, t):
        super(Holder, self).__init__()
        self.t = t
        self.gpudata = t.data_ptr()
    def get_pointer():
        return self.t.data_ptr()

kernel_code = r""" void __global__ add(const float *x, const float *y, float *z) { const int n = threadIdx.y*blockDim.x+threadIdx.x; z[n] = x[n] + y[n]; } """
mod = SourceModule(kernel_code)
add= mod.get_function("add")
A = torch.zeros((10, 10), device="cuda:0")
B = torch.ones((10, 10), device="cuda:0")
C = torch.rand((10, 10), device="cuda:0")

add(Holder(A), Holder(B), Holder(C), grid=(1,1,1), block=(10,10,1))
print(C)

输出:

tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]], device='cuda:0')

注意第四行这句话如果没有的话会报错:

Traceback (most recent call last):
  File "/data/zhuyinghao/jarvis-dain/frame_rate_up_conversion/pycuda_ex2.py", line 27, in <module>
    add(Holder(A), Holder(B), Holder(C), grid=(1,1,1), block=(10,10,1))
  File "/root/anaconda3/lib/python3.7/site-packages/pycuda/driver.py", line 436, in function_call
    func._set_block_shape(*block)
pycuda._driver.LogicError: cuFuncSetBlockShape failed: invalid resource handle

3.参考文献

[1] interop-between-pycuda-and-pytorch
[2] pyCUDA加速你的python代码

今天的文章pycuda pytorch交互——使用pycuda处理pytorch tensor分享到此就结束了,感谢您的阅读,如果确实帮到您,您可以动动手指转发给其他人。

版权声明:本文内容由互联网用户自发贡献,该文观点仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 举报,一经查实,本站将立刻删除。
如需转载请保留出处:https://bianchenghao.cn/28028.html

(0)
编程小号编程小号

相关推荐

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注