| 25 | namespace { |
| 26 | |
| 27 | Status SoftmaxGrad(const Scope& scope, const Operation& op, |
| 28 | const std::vector<Output>& grad_inputs, |
| 29 | std::vector<Output>* grad_outputs) { |
| 30 | // Softmax gradient function. |
| 31 | // p = softmax(x) maps from [batch, n] to [batch, m] |
| 32 | // dp/dx = [dp0/dx0 ... dp0/dxn-1 ] |
| 33 | // [ ... ... ] |
| 34 | // [dpm-1/dx0 ... dpm-1/dxn-1] |
| 35 | // dL/dx = dp/dx * dL/dy |
| 36 | // |
| 37 | // Using alternative formula: |
| 38 | // dL/dx = dL/dy * y - sum(dL/dy * y) * y |
| 39 | // = (dL/dy - sum(dL/dy * y)) * y |
| 40 | auto y = op.output(0); |
| 41 | auto dyy = Mul(scope, grad_inputs[0], y); |
| 42 | auto sum = Reshape(scope, Sum(scope, dyy, {1}), {-1, 1}); |
| 43 | auto sub = Sub(scope, grad_inputs[0], sum); |
| 44 | auto dx = Mul(scope, sub, y); |
| 45 | grad_outputs->push_back(dx); |
| 46 | return scope.status(); |
| 47 | } |
| 48 | REGISTER_GRADIENT_OP("Softmax", SoftmaxGrad); |
| 49 | |
| 50 | bool IsZero(const Scope& scope, const Output& grad) { |