最近在琢磨MobileNet,仍然基于古老的caffe,没有迁移到tensorflow或caffe2上。
caffe存在的一个问题是内存管理较糟糕。对于MobileNet训练,经发现,convolution layer的workspace占用的内存不多,但是batch norm layer和prelu layer(我没用relu,而是用了prelu)占用的内存太多了。因此,修改了caffe的prelu layer和batch norm layer。此外对depthwise separable convolution也进行了一点改进。
1. 修改caffe中的prelu layer。
在in-place计算时,不再缓存bottom data,而是用1个字节缓存bottom data每个数据的符号。
这样每个数字从4字节浮点数变成了1个字节,gpu内存占用量减少到原先的1/4。
2. 改进caffe中的batch norm layer。
参考nvidia caffe的CuDNNBatchNormLayer实现了一个使用cudnn的batch norm layer。
初步测试发现cudnn的batch norm函数可以使用in-place计算,即forward时bottom_data = top_data,backward时foward_diff = backward_diff。
因此在in-place计算时,batch norm layer内部只需要缓存bottom data,比bvlc caffe原先的实现或nvidia caffe的实现都减少一半内存占用。
3. 基于github/yonghenglh6的代码,改进convolution depthwise separable layer。
cudnn7.0中虽然加入了group convolution,但初步测试速度较慢,因此暂时不采用cudnn的group convolution来实现depthwise separable convolution。
yonghenglh6的forward计算效率很高,但backward weight时,由于并行度不够,速度较慢。
因此修改ConvDWSeparableBackwardWeight函数,使CUDA启动n * c * k_h * k_w个并行任务,而不是c * k_h * k_w个并行任务,观察到在训练MobileNet时,对后面若干层()的conv dw layer的backward有明显提速。