add dropout

Author: mcabbott

Project.toml

@@ -7,6 +7,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 

src/NNlib.jl

@@ -40,6 +40,9 @@ for f in ACTIVATIONS
 end
 export sigmoid, hardsigmoid, logsigmoid, thresholdrelu # Aliases
 
+include("dropout.jl")
+export dropout, dropout!
+
 include("softmax.jl")
 export softmax, softmax!, ∇softmax, ∇softmax!, logsoftmax, 
     logsoftmax!, ∇logsoftmax, ∇logsoftmax!, logsumexp

src/dropout.jl

@@ -0,0 +1,152 @@
+using Random, ChainRulesCore
+
+"""
+    dropout([rng], A, p; dims=:)
+
+Returns an array in which each element of `A` is either replaced with zero,
+with probability `p`, or else multiplied by `1/(1-p)`.
+
+By default every element is treated independently.
+With `dims=1`, a choice is made for every value of the 1st index
+i.e. each row of a matrix is either zero or not.
+
+Optional first argument is the random number generator used.
+
+# Examples
+```
+julia> dropout(ones(2, 10), 1/5)
+2×10 Matrix{Float64}:
+ 1.25  1.25  0.0   1.25  1.25  1.25  1.25  1.25  1.25  1.25
+ 1.25  1.25  1.25  0.0   1.25  1.25  0.0   1.25  1.25  1.25
+
+julia> mean(dropout(ones(10^4, 5), 0.3), dims=1)
+1×5 Matrix{Float64}:
+ 0.996  1.00171  1.00629  0.998714  0.992429
+
+julia> dropout(ones(5, 5), 0.7, dims=1)  # whole row the same
+5×5 Matrix{Float64}:
+ 3.33333  3.33333  3.33333  3.33333  3.33333
+ 0.0      0.0      0.0      0.0      0.0
+ 0.0      0.0      0.0      0.0      0.0
+ 3.33333  3.33333  3.33333  3.33333  3.33333
+ 0.0      0.0      0.0      0.0      0.0
+
+julia> mean(dropout(ones(10^4, 5), 0.3, dims=1), dims=1)
+1×5 Matrix{Float64}:
+ 1.00571  1.00571  1.00571  1.00571  1.00571
+```
+"""
+dropout(A::AbstractArray, p::Real; dims = :) = dropout(_rng_from_array(A), A, p; dims)
+
+function dropout(rng::AbstractRNG, A::AbstractArray, p::Real; dims = :)
+    T = float(eltype(A))
+    0 <= p <= 1 || throw(ArgumentError("dropout expects a probability 0 <= p <= 1"))
+    if p > 0
+        dst = similar(A, T)
+        pT = convert(real(T), p)
+        _dropout!(rng, dst, A, pT, dims)
+    else
+        # Not so sure we want fast paths... this tries but doesn't guarantee type-stability,
+        # and the rrule does not have such a fast paths.
+        convert(AbstractArray{T}, A)
+    end
+end
+
+"""
+    dropout!(B, A, p; dims=:)
+
+This does exactly `B .= dropout(A, p; dims)`,
+or rather, it's the implementation of out-of-place [`dropout`](@ref).
+"""
+function dropout!(dst::AbstractArray, src::AbstractArray, p::Real; dims=:)
+    size(dst) == size(src) || throw(DimensionMismatch("dropout! expects output array the same size as input"))
+    0 <= p <= 1 || throw(ArgumentError("dropout expects a probability 0 <= p <= 1"))
+    if p > 0
+        rng = _rng_from_array(A)
+        pT = convert(real(eltype(dst)), p)
+        _dropout!(rng, dst, src, pT, dims)
+    else
+        copyto!(dst, src)
+    end
+end
+
+# This is the easy case in that we can safely use the output array for random numbers.
+function _dropout!(rng::AbstractRNG, dst::AbstractArray, src::AbstractArray, p::Real, dims::Colon)
+    val = convert(eltype(dst), 1/(1-p))
+    rand!(rng, dst)
+    # dst .= (dst.>p) .* val .* src  # hits a SIMD bug
+    _fast_broadcast!(dst, src) do q, x
+        (q>p) * val * x
+    end
+    dst
+end
+
+# For other dims, we we do need to allocate something.
+function _dropout!(rng::AbstractRNG, dst::AbstractArray, src::AbstractArray, p::Real, dims)
+    tmp = similar(dst, ntuple(d -> d in dims ? size(src,d) : 1, ndims(src)))
+    rand!(rng, tmp)
+    val = convert(eltype(dst), 1/(1-p))
+    # One-pass strategy:
+    # dst .= (tmp.>p) .* val .* src
+    # Two-pass strategy:
+    _fast_broadcast!(tmp) do q
+        (q>p) * val
+    end
+    dst .= tmp .* src
+end
+
+# The gradient needs to keep the random choices made, thus store at least a BitArray,
+# but the following way turns out to be faster & simpler:
+function ChainRulesCore.rrule(::typeof(dropout), rng::AbstractRNG, A::AbstractArray, p::Real; dims = :)
+    T = float(eltype(A))
+    val = convert(T, 1/(1-p))
+    keep = if dims isa Colon
+        similar(A, T) 
+    else
+        similar(A, T, ntuple(d -> d in dims ? size(A,d) : 1, ndims(A)))
+    end
+    rand!(rng, keep)
+    Y = @. (keep>p) * A * val
+    function dropout_back(Δ)
+        dY = unthunk(Δ)
+        dA = @. (keep>p) * dY * val
+        (NoTangent(), NoTangent(), dA, NoTangent())
+    end
+    return Y, dropout_back
+end
+
+"""
+    _fast_broadcast!(f, x, y, z...)
+
+This does `x .= f.(x, y, z...)`, but works around
+an issue with broadcasting that prevents SIMD in such cases.
+Can be removed once https://github.com/JuliaLang/julia/issues/43153 is fixed.
+
+Not intended for general use. Does not check sizes!
+"""
+function _fast_broadcast!(f::F, x::Array, yz...) where {F<:Function}
+    bc = Broadcast.instantiate(Broadcast.broadcasted(f, x, yz...))
+    @simd ivdep for I in eachindex(bc)
+        @inbounds x[I] = bc[I]
+    end
+    return x
+end
+function _fast_broadcast!(f::F, x::AbstractArray, yz...) where {F<:Function}
+    # CUDA does not suffer from this bug
+    broadcast!(f, x, x, yz...)
+end
+
+
+"""
+    _rng_from_array(x)
+
+Return the random number generator most appropriate for `x`:
+`CUDA.default_rng()` for `CuArray`,
+else `Random.default_rng()`
+"""
+_rng_from_array(::AbstractArray) = Random.default_rng()
+# _rng_from_array(::CuArray) = CUDA.default_rng()
+
+@non_differentiable _rng_from_array(::Any)
+
+

test/dropout.jl

@@ -0,0 +1,42 @@
+using NNlib, Test, Statistics, Random
+using Zygote, StableRNGs, ChainRulesCore
+
+@testset "dropout" begin
+    # Basics
+    x1 = randn(Float32, 3, 4)
+    @test size(@inferred dropout(x1, 0.1)) == (3, 4)
+    @test size(@inferred dropout(x1, 0.2; dims=2)) == (3, 4)
+    @test size(@inferred dropout(x1, 0.3; dims=(1,2))) == (3, 4)
+    @test eltype(dropout(x1, 0.1)) == Float32
+    @test eltype(dropout(x1, 0.1; dims=1)) == Float32
+    @test eltype(dropout(x1, 0.1; dims=(1,2))) == Float32
+
+    rng =  Random.default_rng()
+    @test size(@inferred dropout(rng, x1, 0.1)) == (3, 4)
+    @test size(@inferred dropout(rng, x1, 0.1; dims=2)) == (3, 4)
+
+    # Values
+    d45 = dropout(trues(100, 100, 100), 0.45)
+    @test mean(d45) ≈ 1 atol=1e-2
+    dpi2 = dropout(fill(pi, 1000), 0.2)
+    @test sort(unique(dpi2)) ≈ [0, 5pi/4]
+    d33 = dropout(fill(3, 10, 1000), 0.3, dims=2)
+    @test sort(unique(vec(d33))) ≈ [0, 3/(1-0.3)]
+
+    # Gradient rule
+    y, back = rrule(dropout, rng, hcat(trues(1000), falses(1000)), 0.45)
+    dx = back(fill(3, 1000, 2))[3]
+    @test !all(iszero, dx[:,2])  # this is why we save the random choices
+    @test sort(unique(vec(dx))) ≈ [0, 3/(1-0.45)]
+
+    @testset "Zygote" begin
+        @test Zygote.gradient(x -> sum(dropout(x, 0.3)), x1)[1] isa Matrix{Float32}
+        @test Zygote.gradient(x -> sum(dropout(rng, x, 0.3)), x1)[1] isa Matrix{Float32}
+        @test Zygote.gradient(x -> sum(dropout(x, 0.3, dims=1)), x1)[1] isa Matrix{Float32}
+
+        f1(x) = sum(dropout(x, 0.5))
+        @test_broken Zygote.hessian(f1, [1.0,2.0,3.0]) == zeros(3, 3)  # forward over reverse
+        @test Zygote.hessian_reverse(f1, [1.0,2.0,3.0]) == zeros(3, 3)
+    end
+end
+

test/runtests.jl

@@ -52,6 +52,10 @@ include("test_utils.jl")
         include("ctc.jl")
     end
 
+    @testset "Dropout" begin
+        include("dropout.jl")
+    end
+
     @testset "Fold/Unfold" begin
         include("fold.jl")
     end