是否可以为混合精度的NN设置GPU亲和力,而FP32和FP16分别用于不同的GPU?

时间:2019-02-15 02:00:50

标签: keras tensorflow nvidia

我有一个GTX 1080和一个RTX2080。我想同时使用它们进行训练,但是由于RTX可以以两倍快的速度处理FP16,所以我想对其进行设置,以便训练使用多GPU和RTX处理FP16层,而GTX处理FP32层。

在tensorflow,pytorch或keras下可能吗?

1 个答案:

答案 0 :(得分:1)

Tensorflow

在TF中,可以为要在其上执行设备的每一层(GPU,CPU或特定GPU,如果您有多个GPU ...)指定。这是使用with tf.device('device_name')语句完成的(您需要提供有意义的device_name)。请参阅Using multiple GPUs部分。

Keras

由于这在TF中是可能的,这意味着,如果您将TF用作Keras后端,那么您也可以在Keras中使用它(Keras只是高级神经网络API)。

请注意,在Keras中,Keras中有一个multi_gpu_model()函数,但是只能在多个GPU上复制整个模型,您不能指定要放置在特定GPU上的层。