TensorFlowのグラジエントを特定のGPUに固定するにはどうすればよいですか？

Question

私は埋め込みレイヤーを持つモデルを作成しました.LSTMはgpu：0に、softmaxレイヤーはgpu：1に作成しました。

そうように私は、トレーニングOPを作成：

optim = tf.train.GradientDescentOptimizer(LEARNING_RATE) 
train_op = optim.minimize(final_loss) 

トレーニングはメモリ不足エラーで失敗します。デバイスの配置ロギングを有効にすると、softmaxの勾配がgpu：1ではなくgpu：0に設定されていることがわかります。 TensorFlowを正しく配置するにはどうすればいいですか？

Answer 1

元の（順方向パス）演算子と同じデバイスにグラデーションを配置するには、オプション引数colocate_gradients_with_ops=Trueをoptim.minimize()に渡す必要があります。あなたのコードは次のようになります：

optim = tf.train.GradientDescentOptimizer(LEARNING_RATE) 
train_op = optim.minimize(final_loss, colocate_gradients_with_ops=True)