サブグラフから新しいオペレーションをアセンブルし、新しくアセンブルされたオペレーションにグラデーションを適用するには

Question

サブグラフ（複数の接続されたオペレーションノードで構成）から新しいオペレーションをアセンブルします。新しい操作に自己設計勾配を適用します。 （ポイントはサブグラフの勾配フローを無視し、NEW OPの出力テンソルからNEW OPの入力テンソルへの勾配をブリッジすることです）。誰か助けてくれることを願っています！

Answer 1

あなたはTensorFlow機能にあなたの部分グラフをラップし、それが働いた

python/framework/function_test.py

@function.Defun(dtype, dtype, dtype) 
    def XentLossGrad(logits, labels, dloss): 
     dlogits = array_ops.reshape(dloss, [-1, 1]) * (
      nn_ops.softmax(logits) - labels) 
     dlabels = array_ops.zeros_like(labels) 
     # Takes exp(dlogits) to differentiate it from the "correct" gradient. 
     return math_ops.exp(dlogits), dlabels 

    @function.Defun(dtype, dtype, grad_func=XentLossGrad) 
    def XentLoss(logits, labels): 
     return math_ops.reduce_sum(labels * math_ops.log(nn_ops.softmax(logits)), 
           1)