make dims as field of Dropout

src/layers/normalise.jl

@@ -13,22 +13,24 @@ end
 _testmode!(m, test) = nothing
 
 """
-    Dropout(p)
+    Dropout(p, dims = :)
 
 A Dropout layer. For each input, either sets that input to `0` (with probability
-`p`) or scales it by `1/(1-p)`. This is used as a regularisation, i.e. it
-reduces overfitting during training.
+`p`) or scales it by `1/(1-p)`. The `dims` argument is to specified the unbroadcasted
+ dimensions, i.e. `dims=1` does dropout along columns and `dims=2` along rows. This is
+ used as a regularisation, i.e. it reduces overfitting during training. see also [`dropout`](@ref).
 
 Does nothing to the input once in [`testmode!`](@ref).
 """
 mutable struct Dropout{F}
   p::F
+  dims::Union{Colon, Int, NTuple{N, Int} where N}
   active::Bool
 end
 
-function Dropout(p)
+function Dropout(p; dims = :)
   @assert 0 ≤ p ≤ 1
-  Dropout{typeof(p)}(p, true)
+  Dropout{typeof(p)}(p, dims, true)
 end
 
 _dropout_shape(s, ::Colon) = size(s)
@@ -36,14 +38,27 @@ _dropout_shape(s, dims) = tuple((i ∉ dims ? 1 : si for (i, si) ∈ enumerate(s
 
 _dropout_kernel(y::T, p, q) where {T} = y > p ? T(1 / q) : T(0)
 
-function (a::Dropout)(x; dims = :)
-  a.active || return x
+
+"""
+    dropout(x, p; dims = :)
+
+The dropout function. For each input, either sets that input to `0` (with probability
+`p`) or scales it by `1/(1-p)`. The `dims` argument is to specified the unbroadcasted
+ dimensions, i.e. `dims=1` does dropout along columns and `dims=2` along rows. This is
+ used as a regularisation, i.e. it reduces overfitting during training.
+"""
+function dropout(x, p; dims = :)
   y = similar(x, _dropout_shape(x, dims))
   rand!(y)
-  y .= _dropout_kernel.(y, a.p, 1 - a.p)
+  y .= _dropout_kernel.(y, p, 1 - p)
   return x .* y
 end
 
+function (a::Dropout)(x)
+  a.active || return x
+  return dropout(x, a.p; dims = a.dims)
+end
+
 _testmode!(a::Dropout, test) = (a.active = !test)
 
 """

test/layers/normalisation.jl

@@ -28,11 +28,12 @@ using Flux.Tracker: data
   @test count(a->a == 0, y) == 0
 
   x = rand(100, 50)
-  m = Dropout(0.5)
-  y = m(x; dims=2)
+  m = Dropout(0.5, dims = 2)
+  y = m(x)
   c = map(i->count(a->a==0, @view y[i, :]), 1:100)
   @test minimum(c) == maximum(c)
-  y = m(x; dims=1)
+  m = Dropout(0.5, dims = 1)
+  y = m(x)
   c = map(i->count(a->a==0, @view y[:, i]), 1:50)
   @test minimum(c) == maximum(c)
 end