如何使用 GPU 训练包含嵌入层的深度学习神经网络？答案

【问题标题】：How do I train deep learning neural network that contains embedding layer using GPU?如何使用 GPU 训练包含嵌入层的深度学习神经网络？
【发布时间】：2021-05-26 02:26:20
【问题描述】：

我的嵌入层出现 InvalidArgumentError：

Colocation Debug Info:
Colocation group had the following types and supported devices: 
Root Member(assigned_device_name_index_=2 requested_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' assigned_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' resource_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' supported_device_types_=[CPU] possible_devices_=[]
GatherV2: GPU CPU 
Cast: GPU CPU 
Const: GPU CPU 
ResourceSparseApplyAdagradV2: CPU 
_Arg: GPU CPU 
ReadVariableOp: GPU CPU 

Colocation members, user-requested devices, and framework assigned devices, if any:
  model_6_user_embedding_embedding_lookup_readvariableop_resource (_Arg)  framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
  adagrad_adagrad_update_1_update_0_resourcesparseapplyadagradv2_accum (_Arg)  framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
  model_6/User-Embedding/embedding_lookup/ReadVariableOp (ReadVariableOp) 
  model_6/User-Embedding/embedding_lookup/axis (Const) 
  model_6/User-Embedding/embedding_lookup (GatherV2) 
  gradient_tape/model_6/User-Embedding/embedding_lookup/Shape (Const) 
  gradient_tape/model_6/User-Embedding/embedding_lookup/Cast (Cast) 
  Adagrad/Adagrad/update_1/update_0/ResourceSparseApplyAdagradV2 (ResourceSparseApplyAdagradV2) /job:localhost/replica:0/task:0/device:GPU:0

     [[{{node model_6/User-Embedding/embedding_lookup/ReadVariableOp}}]] [Op:__inference_train_function_2997]

google colab 的链接： https://colab.research.google.com/drive/1ZN1HzSTTfvA_zstuI-EsKjw7Max1f73v?usp=sharing

这是一个非常简单的神经网络，可以从 Kaggle 下载数据 - 您只需将其拖放到 colabs 中即可使其工作。

我也试过设置soft device placement = True tf.config.set_soft_device_placement(True) 但这似乎不起作用。

从错误日志中，看起来 MirroredStrategy 已将嵌入查找操作分配给 GPU（这是 GPU 不兼容的，我可以看到原因），我希望 tf.config.set_soft_device_placement(True) 会要求 Tensorflow 使用 CPU，但它感觉好像被忽略了。

有没有人以前见过这个问题并知道解决方法？

【问题讨论】：

标签： tensorflow keras gpu embedding

【解决方案1】：

发现TF1.14 的类似问题： https://github.com/tensorflow/tensorflow/issues/31318

看起来 MirroredStrategy 不支持使用基于动量的优化器来训练嵌入层。

克隆上述笔记本并使用 RMSprop（动量=0）似乎可行： https://colab.research.google.com/drive/13MXa8Q96M6uzlkK3K_M7vmQfclL59eRj?usp=sharing

我将使用没有动量的 RMSProp现在，直到这个问题得到解决。错误消息肯定没有帮助！

【讨论】：