当机器上有多个显卡时,TensorFlow如何最大限度的使用呢?目前TensorFlow并没有完美支持多显卡并行,它提倡的是数据并行而已。所以我们在买显卡的时候,尽可能选择单显卡的显存大一些的。
那么多显卡该如何实现数据并行呢?使用的就是下面的API
tensorflow.distribute.Strategy
是一个可在多个 GPU、多台机器或 TPU 上进行分布式训练的 TensorFlow API。使用此 API,您只需改动较少代码就能分布现有模型和训练代码。
tensorflow.distribute.Strategy具有一下特点:
- 使用方便, 支持多用户
- 提供导入即可用的良好性能
- 轻松切换策略
导入该程序后,默认选择你机器上的所有显卡进行训练:
import tensorflow as tf
mirrored_strategy = tf.distribute.MirroredStrategy()
如果你想选择机器上的某几个显卡的话(假设你有5张显卡,你想用第2,3张):
mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:1", "/gpu:2"])
接下来正常写你的代码即可,然后执行,就使用所有显卡进行训练了。
注意:
虽然tf.distribute.MirroredStrategy 支持在一台机器的多个 GPU 上进行同步分布式训练。但是该策略会为每个GPU设备创建一个数据备份,模型中的每个变量都会在副本之间进行复制(镜像)。这些变量会有一个共同的名字,MirroredVariable镜像变量。折现变量会应用相同的更新彼此保持同步。这样做只是增加了计算速度,让你的结果更快的计算出来,但是并不能想CPU并行那样,内存里的变量会共享。也就是说,这里的显卡并行计算,只是提高的速度,并不会使得你的数据量翻倍。如果增加数据,依然回报内存溢出的错误。