GPU + CPU Tensorflow训练

Question

设置

我有一个网络，它的参数是一个大嵌入矩阵(300万×300大小)，embed_mat说。

在训练期间，对于每一小批，我只更新一个小子集的向量从embed_mat (最大15000向量)被选择使用embedding_lookup op。我正在使用Adam优化器来训练我的模型。

由于无法将此embed_mat存储在GPU中，由于它的大小，我在CPU (例如/cpu:0)设备下定义了它，但是模型的其余参数、优化器等都是在GPU (例如gpu:/0)设备下定义的。

问题

1. 我发现我的GPU使用量非常小(200 MB)，这意味着我所有的培训都是在CPU上进行的。我所期待的是，embedding_lookup的结果被复制到GPU，我所有的训练都发生在那里。我是不是做错什么了。
2. 训练时间很大程度上受嵌入矩阵的大小(num_vectors)的影响，这在我看来是不正确的。在任何小型批处理中，我只更新我的网络参数和我查找到的向量(~15000)，因此，如果有的话，训练时间应该与嵌入矩阵的大小呈次线性增长。
3. 有没有一种方法可以自动和无缝地将我的embed_mat分割成多个GPU，以便进行更快的培训？
4. 我怀疑亚当优化器就是因为这个。看起来，因为embed_mat在CPU上，所以所有的训练都在CPU上进行。这是正确的吗？

Answer 1

试着在张紧板上想象一下，你的每个操作都放在哪里。在“图形”选项卡中，你可以用“设备”来着色。理想情况下，嵌入变量、嵌入查找和嵌入梯度更新应该在CPU中，而大多数其他东西应该在GPU中。