TensorFlow的计算平台也分为两种不同的模式,一种是单机训练,另一种是多机分布式训练。
对单机训练来说,虽然执行过程中也包括CPU、GPU的并行计算过程,但总体上处于共享内存的环境,不用过多考虑通信的问题;
而多机分布式训练指的是多台不共享内存的独立节点组成的集群环境下的训练方法。计算节点间需要依靠网络通信,因此可以认为这是与Parameter Server相似的计算环境。
TensorFlow的单机训练是在一个worker节点上进行的,单worker节点内部按照任务关系图的方式在不同GPU+CPU节点间进行计算;对分布式环境来说,平台存在多worker节点,如果采用TensorFlow的Parameter Server策略(tf.disrtibute.experimental.ParameterServerStrategy),则各worker节点会以数据并行的方式进行训练。也就是说,各worker节点以同样的任务关系图的方式进行训练,但训练数据不通,产生的梯度以Parameter Server的方式汇总更新。
接下来介绍每个worker节点内部CPU和GPU的具体任务分工。GPU拥有多核优势,因此在处理矩阵加、向量乘等张量运算时,相比于CPU拥有巨大优势。在处理一个任务节点或任务子图时,CPU主要负责数据和任务的调度,而GPU则负责计算密集度高的张量运算。
举例来说,在处理两个向量的元素乘操作时,CPU会居中调度,把两个向量对应范围的元素发送给GPU处理,再手机处理结果,最终生成处理好的结果向量。从这个角度讲,CPU+GPU的组合也像是一个简化版的“Parameter Server”。
参考资料:
深度学习推荐系统 王喆编著 中国工信出版集团