Add CUDA tests for dropout (#60)

* add dropout

* bump

Author: mcabbott

ext/NNlibCUDA/Project.toml

@@ -1,6 +1,6 @@
 name = "NNlibCUDA"
 uuid = "a00861dc-f156-4864-bf3c-e6376f28a68d"
-version = "0.2.5"
+version = "0.2.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -13,14 +13,15 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 Adapt = "3.3"
 CUDA = "3.11"
-NNlib = "0.8.11"
+NNlib = "0.8.14"
 julia = "1.6"
 
 [extras]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["CUDA", "ForwardDiff", "Test", "Zygote"]
+test = ["CUDA", "ChainRulesCore", "ForwardDiff", "Test", "Zygote"]

ext/NNlibCUDA/src/utils.jl

@@ -1,3 +1,5 @@
+NNlib._rng_from_array(::CuArray) = CUDA.default_rng()
+
 function divide_kernel!(xs, ys, max_idx)
     index = threadIdx().x + (blockIdx().x - 1) * blockDim().x
 

ext/NNlibCUDA/test/dropout.jl

@@ -0,0 +1,34 @@
+using NNlib, NNlibCUDA, CUDA, Test
+using Zygote, ChainRulesCore
+
+@testset "dropout + CUDA" begin
+    # Basics
+    x1 = CUDA.randn(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)
+
+    rng =  CUDA.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(CUDA.ones(100, 100, 100), 0.45)
+    @test mean(d45) ≈ 1 atol=1e-2
+    dpi2 = dropout(CUDA.fill(1f0 * pi, 1000), 0.2)
+    @test sort(unique(Array(dpi2))) ≈ [0, 5pi/4]
+    d33 = dropout(CUDA.fill(3f0, 10, 1000), 0.3, dims=2)
+    @test sort(unique(vec(Array(d33)))) ≈ [0, 3/(1-0.3)]
+
+    # Gradient rule
+    y, back = rrule(dropout, rng, hcat(CUDA.ones(1000), CUDA.zeros(1000)), 0.45)
+    dx = back(CUDA.fill(3f0, 1000, 2))[3]
+    @test !all(iszero, dx[:,2])  # this is why we save the random choices
+    @test sort(unique(vec(Array(dx)))) ≈ [0, 3/(1-0.45)]
+
+    @testset "Zygote" begin
+        @test Zygote.gradient(x -> sum(dropout(x, 0.3)), x1)[1] isa CuArray{Float32}
+        @test Zygote.gradient(x -> sum(dropout(rng, x, 0.3)), x1)[1] isa CuArray{Float32}
+        @test Zygote.gradient(x -> sum(dropout(x, 0.3, dims=1)), x1)[1] isa CuArray{Float32}
+    end
+end

ext/NNlibCUDA/test/runtests.jl

@@ -10,6 +10,7 @@ CUDA.allowscalar(false)
 @testset "NNlibCUDA" begin
 include("test_utils.jl")
 include("activations.jl")
+include("dropout.jl")
 include("batchedadjtrans.jl")
 include("batchedmul.jl")
 include("upsample.jl")