pytorch 的 CUDA 编程 CUDAExtension

版权声明:本文为博主原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。
本文链接:https://blog.csdn.net/lianbus/article/details/94014302

通常我们在使用pytorch设计和训练网络后,需要对网络输出的结果再做一定后处理,才能得到我们的想要的结果。

通常的方法是将网络输出的结果,复制回CPU,并以numpy形式进行数据的后处理,但是对于任务很重的后处理,比如对场数据或者图像数据,做这些操作就会耗费大量的时间,更理想的做法是使用pytorch提供的CUDA扩展,直接在GPU中处理好这些数据,只将最终的结果返回给CPU,一方面并行计算可以极大的缩短后处理的运行时间,另一方面,也节省了数据从GPU直接复制到CPU的时间消耗。

 

教程 https://pytorch.org/docs/master/cpp_extension.html#

关键词 CUDAExtension PYBIND11_MODULE

展开阅读全文

没有更多推荐了,返回首页