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

Question

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

在tensorflow，pytorch或keras下可能吗？

Answer 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上的层。