Merge 6be0b27 into 9a44198

SciML · Sep 6, 2020 · 7080da3 · 7080da3
2 parents 9a44198 + 6be0b27
commit 7080da3
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 9 deletions.
diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
@@ -6,7 +6,7 @@ variables:
 
 .base:
   extends:
-    - .julia:1.4
+    - .julia:1.5
     - .test
   tags:
     - nvidia

diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml
@@ -1,5 +1,5 @@
 [deps]
-CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 MLDataUtils = "cc2ba9b6-d476-5e6d-8eaf-a92d5412d41d"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
diff --git a/test/mnist_gpu.jl b/test/mnist_gpu.jl
@@ -1,8 +1,8 @@
 using DiffEqFlux, OrdinaryDiffEq, Flux, MLDataUtils, NNlib
 using Flux: logitcrossentropy
 using MLDatasets: MNIST
-using CuArrays
-CuArrays.allowscalar(false)
+using CUDA, Test
+CUDA.allowscalar(false)
 
 function loadmnist(batchsize = bs)
 	# Use MLDataUtils LabelEnc for natural onehot conversion
@@ -97,3 +97,4 @@ end
 
 # Train the NN-ODE and monitor the loss and weights.
 Flux.train!( loss, params( down, nn_ode.p, fc), zip( x_train, y_train ), opt, cb = cb )
+@test accuracy(m, zip(x_train,y_train)) > 0.8
diff --git a/test/neural_de_gpu.jl b/test/neural_de_gpu.jl
@@ -1,11 +1,12 @@
-using OrdinaryDiffEq, StochasticDiffEq, Flux, DiffEqSensitivity, DiffEqFlux, Zygote, Test, CuArrays
-CuArrays.allowscalar(false)
+using OrdinaryDiffEq, StochasticDiffEq, Flux, DiffEqSensitivity, DiffEqFlux, Zygote, Test, CUDA
+CUDA.allowscalar(false)
 
 mp = Chain(Dense(2,2)) |> gpu
 x = Float32[2.; 0.] |> gpu
 xs = Float32.(hcat([0.; 0.], [1.; 0.], [2.; 0.])) |> gpu
 tspan = (0.0f0,25.0f0)
 dudt = Chain(Dense(2,50,tanh),Dense(50,2)) |> gpu
+dudt_tracker = Chain(Dense(2,50,CUDA.tanh),Dense(50,2)) |> gpu
 
 NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)(x)
 NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)(x)
@@ -24,7 +25,7 @@ grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
 @test ! iszero(grads[xs])
 @test ! iszero(grads[node.p])
 
-node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
+node = NeuralODE(dudt_tracker,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
 grads = Zygote.gradient(()->sum(Array(node(x))),Flux.params(x,node))
 @test ! iszero(grads[x])
 @test ! iszero(grads[node.p])

diff --git a/test/odenet_gpu.jl b/test/odenet_gpu.jl
@@ -1,8 +1,8 @@
-using DiffEqFlux, Flux, Test, OrdinaryDiffEq, CuArrays
+using DiffEqFlux, Flux, Test, OrdinaryDiffEq, CUDA
 using Statistics
 #= using Plots =#
 
-CuArrays.allowscalar(false)
+CUDA.allowscalar(false)
 
 ## True Solution
 u0 = [2.; 0.] |> gpu

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -57,4 +57,5 @@ if !is_APPVEYOR && GROUP == "GPU"
     end
     @safetestset "odenet GPU" begin include("odenet_gpu.jl") end
     @safetestset "Neural DE GPU Tests" begin include("neural_de_gpu.jl") end
+    @safetestset "MNIST GPU Tests" begin include("mnist_gpu.jl") end
 end