tensorflow layout optimizer && conv autotune

1.Layout Optimizer

  Tensorflow有几种图优化的方法,其中一种较为重要的是layout optimizer,核心思想是调整tensor的layout(NHWC to NCHW),原因在于在较早的cudnn版本中的API是不支持NHWC格式的输入的,目前cudnn7.0版本已经能支持NHWC格式输入了,但经过实测发现以NHWC格式为输入调用cudnn API的速度是不如NCHW的(在k80上实测,lenet-5的两个卷积层,使用NHWC计算需要平均约2.5ms,使用NCHW只需要1.8ms),所以Tensorflow layout的作用就是将NHWC转换为NCHW,从而加速网络的计算速度。

  在使用layout转换的时候,Tensorflow会监测网络的结构,如果发现存在调用cudnn计算的节点,就会在节点前加一个transpose,将tensor转换为NCHW格式,然后经过所有由cudnn计算的节点及计算方式与layout无关的节点(比如relu的计算就与layout无关,但是reduce_mean就与layout有关)之后,再加一个transpose将tensor转换回来。以resnet_v2为例:https://github.com/tensorpack/tensorpack/blob/master/examples/ResNet/cifar10-resnet.py,resnet的中间网络都是Conv + BN + Relu的组合,因此像这种网络Tensorflow就会在第一个Conv节点前加一个transpose节点,,在GlobalAvgPooling前再加一个transpose节点(因为GlobalAvgPooling调用的是reduce_mean(input, [1,2])),用timeline分析节点运算时间的应该注意到过这两个transpose节点。这样就用两个transpose的时间换取更快的cudnn的调用时间,大幅缩减网络的执行时间。

  但是假设你的网络中调用cudnn的节点不连续,就会造成加多个transpose的情况。虽然也能起到加速效果,但是肯定不如只加两个transpose更合理,这样在设计网路的时候才能设计出运行更快的网络.

2.Conv Autotune

  cudnn的conv forward API:https://docs.nvidia.com/deeplearning/sdk/cudnn-developer-guide/index.html#cudnnConvolutionForward,其中有一个algo参数,含义是计算卷积所用的算法,这个参数一共可以传8种不同的算法,具体参考:https://docs.nvidia.com/deeplearning/sdk/cudnn-developer-guide/index.html#cudnnConvolutionFwdAlgo_t,在真正计算卷积的时候是不能确定哪种算法速度是最快的,因此Tensorflow启用了一种缓存机制,就是对于每种尺寸的Tensor和conv参数(strides, padding啥的),8个算法依次调用。哪次时间短就将其记录下来,下次继续用(其实cudnn有相关的API可以直接调用的,Tensorflow1.13版本,貌似还是自己实现的调用+计时功能),相当于对于每种不同的尺寸的输入都要算8次卷积,所以如果模型是变尺寸的模型,那么这8次卷积带来的代价就是严重拖慢模型的速度。那么是否可以关闭这个缓存功能呢?Tensorflow的环境变量TF_CUDNN_USE_AUTOTUNE可以控制是否打开卷积的autotune功能。那么如果关闭了这个autotune,是否能找到运算最快的卷积算法呢,答案是不能.关闭autotune功能。就只能调用cudnnGetConvolutionForwardAlgorithm函数以一种启发式的搜索方式返回一个算法,这个返回值通常都不是最优解,所以如果网络是变尺寸的情况下,最好的处理方式就是在一定范围内缓存所有情况的卷积输入尺寸与使用算法的对应关系。当然最好的情况还是固定卷积的尺寸.