Wrapper functions for NNlib (#615)

Author: matsueushi

src/nnlib.jl

@@ -1,34 +1,26 @@
 using NNlib
 
 # Activation functions
+@cufunc softplus(x::Real) = ifelse(x > 0, x + log1p(exp(-x)), log1p(exp(x)))
 
-@cufunc σ(x) = ifelse(x < -80, zero(x), one(x) / (one(x) + exp(-x)))
+@cufunc logσ(x::Real) = -softplus(-x)
 
-@cufunc function logσ(x)
-  max_v = max(zero(x), -x)
-  z = exp(-max_v) + exp(-x-max_v)
-  -(max_v + log(z))
+@cufunc function gelu(x::Real)
+    p = oftype(x / 1, π)
+    λ = oftype(x / 1, √(2 / p))
+    α = oftype(x / 1, 0.044715)
+    h = oftype(x / 1, 0.5)
+    h * x * (one(x) + tanh(λ * (x + α * x^3)))
 end
 
-@cufunc elu(x, α = one(x)) =
-  ifelse(x ≥ 0, x/1, α * (exp(x) - one(x)))
+@cufunc lisht(x::Real) = x * tanh(x)
 
-@cufunc swish(x) = x * σ(x)
+@cufunc logcosh(x::Real) = x + softplus(-2x) - log(oftype(x, 2))
 
-@cufunc function gelu(x)
-  λ = oftype(x/1, √(2/π))
-  α = oftype(x/1, 0.044715)
-  h = oftype(x/1, 0.5)
-  h * x * (one(x) + tanh(λ * (x + α * x^3)))
-end
+@cufunc mish(x::Real) = x * tanh(softplus(x))
 
-@cufunc function selu(x)
-  λ = oftype(x/1, 1.0507009873554804934193349852946)
-  α = oftype(x/1, 1.6732632423543772848170429916717)
-  λ * ifelse(x > 0, x/1, α * (exp(x) - 1))
-end
+@cufunc tanhshrink(x::Real) = x - tanh(x)
 
-@cufunc softplus(x) = ifelse(x > 0, x + log1p(exp(-x)), log1p(exp(x)))
 
 
 # Batched matrix multiplication

test/dnn.jl

@@ -80,8 +80,11 @@ end
   @test testf(CUDNN.cudnnActivationBackward, cu(rand(Float64, 10, 10, 3, 1)), cu(rand(Float64, 10, 10, 3, 1)), cu(rand(Float64, 10, 10, 3, 1)), cu(rand(Float64, 10, 10, 3, 1)))
 
   # activations defined in src/nnlib.jl
+  ACTIVATION_FUNCTIONS = [σ, logσ, hardσ, hardtanh, relu, leakyrelu, relu6, rrelu, 
+                          elu, gelu, celu, swish, lisht, selu, trelu, softplus, 
+                          softsign, logcosh, mish, tanhshrink, softshrink];
   for dims in ((5,5), (5,))
-    for f in (σ, logσ, elu, swish, gelu, selu, softplus)
+    for f in filter(x -> x != rrelu, ACTIVATION_FUNCTIONS)
       @test testf(x -> f.(x), rand(Float64, dims))
     end
   end