use NNlib.bias_act

Author: mcabbott

Project.toml

@@ -44,7 +44,7 @@ Functors = "0.4"
 MLUtils = "0.4"
 MacroTools = "0.5"
 Metal = "0.5, 1"
-NNlib = "0.9.1"
+NNlib = "0.9.5"
 OneHotArrays = "0.2.4"
 Optimisers = "0.3.2"
 Preferences = "1"

src/layers/basic.jl

@@ -169,9 +169,8 @@ end
 
 function (a::Dense)(x::AbstractVecOrMat)
   _size_check(a, x, 1 => size(a.weight, 2))
-  σ = NNlib.fast_act(a.σ, x)  # replaces tanh => tanh_fast, etc
   xT = _match_eltype(a, x)  # fixes Float64 input, etc.
-  return σ.(a.weight * xT .+ a.bias)
+  NNlib.bias_act!(a.σ, a.weight * xT, a.bias)  # does σ.(W*x .+ b), with fast paths
 end
 
 function (a::Dense)(x::AbstractArray)
@@ -446,7 +445,7 @@ function (a::Bilinear)(x::AbstractMatrix, y::AbstractMatrix)
   Z = reshape(Wyx, (d_z, :))
 
   # @einsum out[o,s] := σ(Z[o,i] + b[o])
-  σ.(Z .+ b)
+  NNlib.bias_act!(σ, Z, b)  # σ.(Z .+ b)
 end
 
 (a::Bilinear)(x::AbstractVecOrMat) = a(x, x)

src/layers/conv.jl

@@ -196,10 +196,11 @@ ChainRulesCore.@non_differentiable conv_dims(::Any, ::Any)
 
 function (c::Conv)(x::AbstractArray)
   _size_check(c, x, ndims(x)-1 => _channels_in(c))
-  σ = NNlib.fast_act(c.σ, x)
+  # σ = NNlib.fast_act(c.σ, x)
   cdims = conv_dims(c, x)
   xT = _match_eltype(c, x)
-  σ.(conv(xT, c.weight, cdims) .+ conv_reshape_bias(c))
+  # σ.(conv(xT, c.weight, cdims) .+ conv_reshape_bias(c))
+  NNlib.bias_act!(c.σ, conv(xT, c.weight, cdims), conv_reshape_bias(c))
 end
 
 _channels_in(l::Conv) = size(l.weight, ndims(l.weight)-1) * l.groups
@@ -332,10 +333,11 @@ ChainRulesCore.@non_differentiable conv_transpose_dims(::Any, ::Any)
 
 function (c::ConvTranspose)(x::AbstractArray)
   _size_check(c, x, ndims(x)-1 => _channels_in(c))
-  σ = NNlib.fast_act(c.σ, x)
+  # σ = NNlib.fast_act(c.σ, x)
   cdims = conv_transpose_dims(c, x)
   xT = _match_eltype(c, x)
-  σ.(∇conv_data(xT, c.weight, cdims) .+ conv_reshape_bias(c))
+  # σ.(∇conv_data(xT, c.weight, cdims) .+ conv_reshape_bias(c))
+  NNlib.bias_act!(c.σ, ∇conv_data(xT, c.weight, cdims), conv_reshape_bias(c))
 end
 
 function Base.show(io::IO, l::ConvTranspose)
@@ -474,10 +476,11 @@ ChainRulesCore.@non_differentiable crosscor_dims(::Any, ::Any)
 
 function (c::CrossCor)(x::AbstractArray)
   _size_check(c, x, ndims(x)-1 => _channels_in(c))
-  σ = NNlib.fast_act(c.σ, x)
+  # σ = NNlib.fast_act(c.σ, x)
   cdims = crosscor_dims(c, x)
   xT = _match_eltype(c, x)
-  σ.(crosscor(xT, c.weight, cdims) .+ conv_reshape_bias(c))
+  # σ.(crosscor(xT, c.weight, cdims) .+ conv_reshape_bias(c))
+  NNlib.bias_act!(c.σ, crosscor(xT, c.weight, cdims), conv_reshape_bias(c))
 end
 
 function Base.show(io::IO, l::CrossCor)

src/layers/normalise.jl

@@ -246,7 +246,7 @@ function _norm_layer_forward(
   β = reshape(l.β, affine_shape)
 
   scale = γ ./ sqrt.(σ² .+ eps)
-  bias = -scale .* μ .+ β
+  bias = .-scale .* μ .+ β
   l.λ.(scale .* x .+ bias)
 end