Tensorflow教程:GPU调用如何实现

今天,给大家分析一下Tensorflow源码之GPU调用是如何实现的?

1. Tensorflow GPU支持

Tensorflow 支持GPU进行运算,目前官方版本只支持NVIDIA的GPU,可以在tensorflow的官方上看到。Tensorflow 对GPU的运算的支持最小力度就是OP,也就是我们常说的算子,下图提供了Tensorflow的一些常见算子,而每个算子在Tensorflow上都会提供GPU的算法:关于OP的具体实现,在本篇博客中就不叙述了。

2. Tensorflow GPU调用架构

从上图我们可以看到,Tensorflow提供两种方式调用NVIDIA的方式,而NVIDIA的GPU调用方式主要依靠的CUDA的并行计算框架

2.1 Stream Executor

StreamExecutor 是一个子项目,是一个google开源的数学并行运算库,是基于CUDA API、OpenCL API管理各种GPU设备的统一API,这种统一的GPU封装适用于需要与GPU设备通信的库,而在Tensorflow上只提供了对CUDA的支持

StreamExecutor的主要功能:

  • 抽象化底层平台,对开发者不需要考虑底层的GPU的平台
  • 流式的管理模式
  • 封装了主机和GPU之间的数据移动

在StreamExecutor里封装了几个常见的基本的核心运算:

  • BLAS: 基本线性代数
  • DNN:  深层神经网络
  • FFT:   快速傅里叶变换
  • RNG:  随机数生成

2.1.1 Stream 接口

  1.  算子直接通过Stream的API的调用,在Tensorflow里Stream executor 只支持4个核心算法
  2.  每个算法都提供Support的类,进行多态的支持,比如CUDA, OpenCL
  3.  通过Support,官方tensorflow 只提供了CUDA支持,如果要支持OpenCL,可以参考开源(点击打开链接
  4.  对CUDA的支持使用了基于CUDA平台的第三方开发库,没有直接使用CUDA编程

2.2  直接调用CUDA

Tensorflow 同时本身也可以直接调用CUDA,毕竟Stream的目前接口只是支持了Blas, DNN, FFT, RND这些基本接口

1.  进行复杂运算,需要连续调用Stream的接口,这里也带来了频繁的从主内存到GPU内存之间复制的开销

2.  Stream 并没有封装一些简单的一元运算,只是封装了CUDA的提供的第三方运算库,一元运算(加减乘除,log, exp)这些如果想在GPU运算,需要基于CUDA的运算框架进行自己写代码

在Tensorflow上写CUDA代码没什么两样, 下面是一个lstm的样例

1. 定义你的global 

[html] view plain copy

print?

  1. template <typename T, bool use_peephole>
  2. __global__ void lstm_gates(const T* icfo, const T* b, const T* cs_prev,  
  3.                            const T* wci, const T* wcf, const T* wco, T* o, T* h,  
  4.                            T* ci, T* cs, T* co, T* i, T* f, const T forget_bias,  
  5.                            const T cell_clip, const int batch_size,  
  6.                            const int cell_size) {  
  7.   const int batch_id = blockIdx.x * blockDim.x + threadIdx.x;  
  8.   const int act_id = blockIdx.y * blockDim.y + threadIdx.y;  
  9. .......  
  10. }  

2. 定义使用的网格,block, thread数

[html] view plain copy

print?

  1. dim3 block_dim_2d(std::min(batch_size, 8), 32);  
  2. dim3 grid_dim_2d(Eigen::divup(batch_size, static_cast<int>(block_dim_2d.x)),  
  3.                  Eigen::divup(cell_size, static_cast<int>(block_dim_2d.y)));  
  4. if (use_peephole) {  
  5.   lstm_gates<T, true><<<grid_dim_2d, block_dim_2d, 0, cu_stream>>>(  
  6.       icfo.data(), b.data(), cs_prev.data(), wci.data(), wcf.data(),  
  7.       wco.data(), o.data(), h.data(), ci.data(), cs.data(), co.data(),  
  8.       i.data(), f.data(), forget_bias, cell_clip, batch_size, cell_size);  
  9. } else {  
  10.   lstm_gates<T, false><<<grid_dim_2d, block_dim_2d, 0, cu_stream>>>(  
  11.       icfo.data(), b.data(), cs_prev.data(), wci.data(), wcf.data(),  
  12.       wco.data(), o.data(), h.data(), ci.data(), cs.data(), co.data(),  
  13.       i.data(), f.data(), forget_bias, cell_clip, batch_size, cell_size);  
  14. }  

3. 定义你的OP,在你的OP里调用CUDA的代码,并注册到Tensorflow Kernel中,注意你的Device需要设置成DEVICE_GPU,tensorflow会依据客户端传递的device的参数来决定是否需调用GPU还是CPU的算法,CUDA的文件以.cu.cc为结尾

[html] view plain copy

print?

  1. REGISTER_KERNEL_BUILDER(  
  2.     Name("arithmetic").Device(DEVICE_GPU).TypeConstraint<Eigen::half>("T"),  
  3.     arithmeticOP<Eigen::half>);  

原创声明,本文系作者授权云+社区发表,未经许可,不得转载。

如有侵权,请联系 yunjia_community@tencent.com 删除。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏Petrichor的专栏

tensorflow: tf.reshape探究

  给定一个tensor,这个操作会返回一个有着跟原tensor一样的值且经过shape重塑过的张量。

991
来自专栏漫漫深度学习路

pytorch学习笔记(二):gradient

gradient 在BP的时候,pytorch是将Variable的梯度放在Variable对象中的,我们随时都可以使用Variable.grad得到对应Var...

2518
来自专栏机器之心

一段Python代码,告诉你机器之心今天的秘密

31411
来自专栏浪淘沙

SparkSql学习笔记一

1.简介     Spark SQL是Spark用来处理结构化数据的一个模块,它提供了一个编程抽象叫做DataFrame并且作为分布式SQL查询引擎的作用。 ...

1113
来自专栏AILearning

使用GPU

支持的设备 在典型的系统中,有多个计算设备。在TensorFlow中,支持的设备类型是CPU和GPU。它们被表示为strings。例如: "/cpu:0":...

1805
来自专栏ml

使用神经网络来拟合函数y = x^3 +b

我们使用一个三层的小网络来,模拟函数y = x^3+b函数 1 import tensorflow as tf 2 import numpy as np ...

30811
来自专栏专知

不只是支持Windows, PyTorch 0.4新版本变动详解与升级指南

4702
来自专栏伦少的博客

Spark UDF使用详解及代码示例

本文介绍如何在Spark Sql和DataFrame中使用UDF,如何利用UDF给一个表或者一个DataFrame根据需求添加几列,并给出了旧版(Spark1....

1483
来自专栏大数据智能实战

Spark 2.0 DataFrame map操作中Unable to find encoder for type stored in a Dataset.问题的分析与解决

随着新版本的spark已经逐渐稳定,最近拟将原有框架升级到spark 2.0。还是比较兴奋的,特别是SQL的速度真的快了许多。。 然而,在其中一个操作时却卡住了...

3909
来自专栏深度学习之tensorflow实战篇

python下的Pandas中DataFrame基本操作(一),基本函数整理

1.6K8

扫码关注云+社区