Gradient function for SoftmaxCrossEntropyWithLogits.
(op, grad_loss, grad_grad)
| 521 | |
| 522 | @ops.RegisterGradient("SoftmaxCrossEntropyWithLogits") |
| 523 | def _SoftmaxCrossEntropyWithLogitsGrad(op, grad_loss, grad_grad): |
| 524 | """Gradient function for SoftmaxCrossEntropyWithLogits.""" |
| 525 | # grad_loss is the backprop for cost, and we multiply it with the gradients |
| 526 | # (which is output[1]) |
| 527 | # grad_grad is the backprop for softmax gradient. |
| 528 | # |
| 529 | # Second derivative is just softmax derivative w.r.t. logits. |
| 530 | softmax_grad = op.outputs[1] |
| 531 | grad = _BroadcastMul(grad_loss, softmax_grad) |
| 532 | |
| 533 | def IsZero(g): |
| 534 | # Some introspection to check if the gradient is feeding zeros |
| 535 | if context.executing_eagerly(): |
| 536 | # TODO(apassos) add an efficient way to detect eager zeros here. |
| 537 | return False |
| 538 | if g.op.type in ("ZerosLike", "Zeros"): |
| 539 | return True |
| 540 | const_fill_value = tensor_util.constant_value(g) |
| 541 | return const_fill_value is not None and (const_fill_value == 0).all() |
| 542 | |
| 543 | logits = op.inputs[0] |
| 544 | if grad_grad is not None and not IsZero(grad_grad): |
| 545 | softmax = nn_ops.softmax(logits) |
| 546 | |
| 547 | grad += ((grad_grad - array_ops.squeeze( |
| 548 | math_ops.matmul( |
| 549 | array_ops.expand_dims(grad_grad, 1), |
| 550 | array_ops.expand_dims(softmax, 2)), |
| 551 | axis=1)) * softmax) |
| 552 | |
| 553 | return grad, _BroadcastMul(grad_loss, -nn_ops.log_softmax(logits)) |
| 554 | |
| 555 | |
| 556 | @ops.RegisterGradient("SparseSoftmaxCrossEntropyWithLogits") |
nothing calls this directly
no test coverage detected