setmode! -> testmode!

src/Flux.jl

@@ -9,7 +9,7 @@ using Lazy: @forward
 
 export Chain, Dense, RNN, LSTM, Dropout,
   SGD, ADAM, Momentum, Nesterov,
-  param, params, mapleaves, setmode!
+  param, params, mapleaves, testmode!
 
 using NNlib
 export σ, relu, leakyrelu, elu, swish, softmax

src/layers/basic.jl

@@ -81,22 +81,22 @@ end
 
 
 """
-  Dropout(p; mode=:train)
+  Dropout(p; testmode=false)
 
-A Dropout layer. In `:train` mode sets input components `x[i]` to zero with
+A Dropout layer. If `testmode=false` mode sets input components `x[i]` to zero with
 probability `p` and to `x[i]/(1-p)` with probability `(1-p)`.
 
-In `:eval` mode it doesn't alter the input: `x == Dropout(p; mode=:eval)(x)`.
-Change the mode with [`setmode!`](@ref).
+In `testmode=true`it doesn't alter the input: `x == Dropout(p; mode=:eval)(x)`.
+Change the mode with [`testmode!`](@ref).
 """
 mutable struct Dropout{F}
   p::F
-  mode::Symbol
+  testmode::Bool
 end
-Dropout(p::F; mode=:train) where {F} = Dropout{F}(p, mode)
+Dropout(p::F; testmode::Bool=false) where {F} = Dropout{F}(p, testmode)
 
 function (a::Dropout)(x)
-  if a.mode == :eval
+  if a.testmode
     return x
   else
     if 0 < a.p < 1
@@ -116,12 +116,9 @@ function (a::Dropout)(x)
 end
 
 """
-    setmode!(m, mode::Symbol)
+    testmode!(m, val=true)
 
-Change the mode of model `m` to `mode`. Possible values for `mode` are
-`:train` and `:eval`.
-This has an affect only if `m` contains [`Dropout`](@ref) of `BatchNorm` layers.
+Set model `m` in test mode if `val=true`, and in training mode otherwise.
+This has an affect only if `m` contains [`Dropout`](@ref) or `BatchNorm` layers.
 """
-setmode!(a, mode::Symbol) = nothing
-setmode!(c::Chain, mode::Symbol) = mapchildren(x->setmode!(x, mode), c)
-setmode!(a::Dropout, mode::Symbol) = a.mode = mode
+testmode!(m, val::Bool=true) = prefor(x -> x isa Dropout && (x.testmode = val), m)

src/tracker/Tracker.jl

@@ -41,6 +41,7 @@ TrackedArray(x::AbstractArray) = TrackedArray(Call(nothing), x, zeros(x))
 param(xs) = TrackedArray(AbstractFloat.(xs))
 istracked(x::TrackedArray) = true
 data(x::TrackedArray) = x.data
+# data(x::TrackedScalar) = x.data[]
 grad(x::TrackedArray) = x.grad
 
 # Fallthrough methods

test/layers.jl

@@ -1,23 +1,26 @@
 @testset "dropout" begin
   x = [1.,2.,3.]
-  @test x === Dropout(0.1, mode=:eval)(x)
-  @test x === Dropout(0, mode=:train)(x)
-  @test all(zeros(x) .== Dropout(1, mode=:train)(x))
+  @test x === Dropout(0.1, testmode=true)(x)
+  @test x === Dropout(0, testmode=false)(x)
+  @test all(zeros(x) .== Dropout(1, testmode=false)(x))
 
   x = rand(100)
   m = Dropout(0.9)
   y = m(x)
   @test count(a->a==0, y) > 50
-  setmode!(m, :eval)
+  testmode!(m)
   y = m(x)
   @test count(a->a==0, y) == 0
+  testmode!(m, false)
+  y = m(x)
+  @test count(a->a==0, y) > 50
 
   x = rand(100)
   m = Chain(Dense(100,100),
             Dropout(0.9))
   y = m(x)
   @test count(a->a == 0, y) > 50
-  setmode!(m, :eval)
+  testmode!(m)
   y = m(x)
   @test count(a->a == 0, y) == 0
 end