pytorch0.4添加自定义cuda层

1. cuda函数的基础类型c或者c++

pytorch扩展工具函数时,c或者c++是分开的。基本的命令和g++有很多相似之处。所以,打算使用不同的扩展方式时,需要相应的源码是c或者c++的。不然会出很多乱七八糟的问题。
另外,pytorch 的官方网站上能找到完全用c语言写forward和backward的例子,我这里是将核心函数由cuda和c来编写,数据流的处理还是python下写的方式来。

2. 使用c的扩展方式

a.参考的博客

有一些教程,可以直接参考下,这一篇oldpan的文章翻译的不错,相比起来,官方的1.0的文档要要与0.4的文档。官方给的一个github版的实例,和oldpan的是同源的,也可以参考下。此外,debug的时候,小郑的数据处理也帮我不少,也记在这。

b.特别留意的点
  1. 输入数据的数据类型
    我是直接读取的图片,是uint8类型的,送入pytorch函数中的时候,很容易忽略转成float类型。
  2. pytorch 的Tensor转成指针
    THCudaTensor_data就好,这儿我下功夫找官方的解释,小郑的里面这里是比较全的。
  3. float3和float类型的转变
    cuda中非常容易发现float234这些类型,数据处理的时候用后者非常的方便,数据拷贝的时候需要注意,通过强制转换指针类型来相互复制。他们在硬件上存储都是一致的,区别在于两个指针的步长不一样。甚至和uchar3的相互转换,也是同样的道理。
    float *dev_image_in,  *image_out;
    float3 *image_in, *dev_image_tmp;
    int memSize = X_SIZE * Y_SIZE * sizeof(float3);
    gpuErrchk(cudaMalloc(&image_in, memSize));
    printf("size float:%ld,size float3:%ld,and uchar3:%ld\n", sizeof(float), sizeof(float3), sizeof(uchar3));
    gpuErrchk(cudaMemcpy((float *) image_in, dev_image_in, memSize, cudaMemcpyHostToDevice));
    //
    memSize = w * h * sizeof(float3);
    gpuErrchk(cudaMemcpy((float3 *) image_out, dev_image_tmp, memSize, cudaMemcpyDeviceToHost));
    
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值