Merge pull request #22 from boathit/master

Adding logsoftmax

src/NNlib.jl

@@ -3,13 +3,14 @@ module NNlib
 using Requires
 
 export σ, sigmoid, relu, leakyrelu, elu, swish, selu, softplus, softsign,
-  softmax, conv2d, maxpool2d, avgpool2d
+  softmax, logsoftmax, conv2d, maxpool2d, avgpool2d
 
 const libnnlib = Libdl.find_library("nnlib.$(Libdl.dlext)", [joinpath(@__DIR__, "..", "deps")])
 
 include("numeric.jl")
 include("activation.jl")
 include("softmax.jl")
+include("logsoftmax.jl")
 include("linalg.jl")
 include("conv.jl")
 

src/logsoftmax.jl

@@ -0,0 +1,34 @@
+using Base.Threads
+
+function logsoftmax!(out::AbstractVecOrMat, xs::AbstractVecOrMat)
+  @threads for j = 1:size(xs, 2)
+    @inbounds begin
+      xi_max = xs[1, j]
+      for i = 1:size(out, 1)
+        xi_max = max(xi_max, xs[i, j])
+      end
+      s = zero(eltype(out))
+      for i = 1:size(out, 1)
+        s += exp(xs[i, j] - xi_max)
+      end
+      for i = 1:size(out, 1)
+        out[i, j] = xs[i, j] - log(s) - xi_max
+      end
+    end
+  end
+  out
+end
+
+
+logsoftmax!(xs) = logsoftmax!(xs, xs)
+logsoftmax(xs) = logsoftmax!(similar(xs), xs)
+
+∇logsoftmax(Δ, xs) = ∇softmax(Δ ./ softmax(xs), xs)
+
+"""
+    logsoftmax(xs) = log.(exp.(xs) ./ sum(exp.(xs)))
+
+logsoftmax computes the log of softmax(xs) and it is more numerically stable
+than softmax function in computing the cross entropy loss.
+"""
+logsoftmax

test/runtests.jl

@@ -8,8 +8,16 @@ include("conv.jl")
 
 xs = rand(5)
 @test softmax(xs) ≈ exp.(xs) ./ sum(exp.(xs))
+@test logsoftmax(xs) ≈ log.(softmax(xs))
 
 xs = rand(5,10)
 @test softmax(xs) ≈ exp.(xs) ./ sum(exp.(xs),1)
+@test logsoftmax(xs) ≈ log.(softmax(xs))
 
+## compare the outputs with the PyTorch nn.LogSoftmax returns
+xs = Float32[1, 2, 3000.]
+@test logsoftmax(xs) ≈ [-2999, -2998, 0]
+
+xs = Float32[1 2 3; 1000 2000 3000]
+@test logsoftmax(xs) ≈ [-999 -1998 -2997; 0 0 0.]
 end