numeric stability

FluxML · Sep 12, 2017 · f446498 · CarloLucibello · Sep 13, 2017 · MikeInnes
1 parent 3fca9d0
commit f446498
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/activation.jl b/src/activation.jl
@@ -2,8 +2,10 @@
 
 relu(x) = max(0, x)
 
-softmax!(out::AbstractVecOrMat, xs::AbstractVecOrMat) =
-  out .= exp.(xs) ./ sum(exp, xs, 1)
+function softmax!(out::AbstractVecOrMat, xs::AbstractVecOrMat)
+  out .= xs .- maximum(xs)
+  out .= exp.(out) ./ sum(exp, out, 1)
+end
 
 softmax!(xs) = softmax!(xs, xs)
 softmax(xs) = softmax!(similar(xs), xs)