basic softmax gradient

FluxML · Aug 22, 2017 · 3fca9d0 · 3fca9d0
1 parent 19cbce5
commit 3fca9d0
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/activation.jl b/src/activation.jl
@@ -6,5 +6,12 @@ softmax!(out::AbstractVecOrMat, xs::AbstractVecOrMat) =
   out .= exp.(xs) ./ sum(exp, xs, 1)
 
 softmax!(xs) = softmax!(xs, xs)
-
 softmax(xs) = softmax!(similar(xs), xs)
+
+function ∇softmax!(out::AbstractVecOrMat, Δ::AbstractVecOrMat, xs::AbstractVecOrMat)
+  s = sum(exp, xs, 1)
+  out .= exp.(xs)./s.*(Δ .- sum(Δ .* exp.(xs), 1)./s)
+end
+
+∇softmax!(Δ, xs) = ∇softmax!(Δ, Δ, xs)
+∇softmax(Δ, xs) = ∇softmax!(similar(Δ), Δ, xs)