Basic GPU Support

examples/GPU/README.md

@@ -0,0 +1,9 @@
+## Testing `RobustNeuralNetworks.jl` on a GPU
+
+There is currently full support for using models from `RobustNeuralNetworks.jl` on a GPU. However, the speed could definitely be improved with some code optimisations, and we don't have any CI testing on the GPU.
+
+The scripts in this directory serve two purposes:
+- They provide a means of benchmarking model performance on a GPU
+- They act as unit tests to verify the models can be trained on a GPU
+
+There is an [open issue](https://github.com/acfr/RobustNeuralNetworks.jl/issues/119) on improving the speed of our models on GPUs. Any and all contributions are welcome.

examples/GPU/test_lbdn.jl

@@ -0,0 +1,59 @@
+# This file is a part of RobustNeuralNetworks.jl. License is MIT: https://github.com/acfr/RobustNeuralNetworks.jl/blob/main/LICENSE 
+
+cd(@__DIR__)
+using Pkg
+Pkg.activate("../")
+
+using BenchmarkTools
+using CUDA
+using Flux
+using Random
+using RobustNeuralNetworks
+
+rng = Xoshiro(42)
+
+function test_lbdn_device(device; nu=2, nh=[10, 5], ny=4, γ=10, nl=tanh, 
+                          batches=4, is_diff=false, do_time=true, T=Float32)
+
+    # Build model
+    model = DenseLBDNParams{T}(nu, nh, ny, γ; nl, rng) |> device
+    is_diff && (model = DiffLBDN(model))
+
+    # Create dummy data
+    us = randn(rng, T, nu, batches) |> device
+    ys = randn(rng, T, ny, batches) |> device
+
+    # Dummy loss function
+    function loss(model, u, y)
+        m = is_diff ? model : LBDN(model)
+        return Flux.mse(m(u), y)
+    end
+
+    # Run and time, running it once to check it works
+    print("Forwards: ")
+    l = loss(model, us, ys)
+    do_time && (@btime $loss($model, $us, $ys))
+
+    print("Reverse:  ")
+    g = gradient(loss, model, us, ys)
+    do_time && (@btime $gradient($loss, $model, $us, $ys))
+
+    return l, g
+end
+
+function test_lbdns(device)
+
+    d = device === cpu ? "CPU" : "GPU"
+    println("\nTesting LBDNs on ", d, ":")
+    println("--------------------\n")
+
+    println("Dense LBDN:\n")
+    test_lbdn_device(device)
+    println("\nDense DiffLBDN:\n")
+    test_lbdn_device(device; is_diff=true)
+
+    return nothing
+end
+
+test_lbdns(cpu)
+test_lbdns(gpu)

examples/GPU/test_ren.jl

@@ -0,0 +1,93 @@
+# This file is a part of RobustNeuralNetworks.jl. License is MIT: https://github.com/acfr/RobustNeuralNetworks.jl/blob/main/LICENSE 
+
+cd(@__DIR__)
+using Pkg
+Pkg.activate("../")
+
+using BenchmarkTools
+using CUDA
+using Flux
+using Random
+using RobustNeuralNetworks
+
+rng = Xoshiro(42)
+
+function test_ren_device(device, construct, args...; nu=4, nx=5, nv=10, ny=4, 
+                         nl=tanh, batches=4, tmax=3, is_diff=false, T=Float32,
+                         do_time=true)
+
+    # Build the ren
+    model = construct{T}(nu, nx, nv, ny, args...; nl, rng) |> device
+    is_diff && (model = DiffREN(model))
+
+    # Create dummy data
+    us = [randn(rng, T, nu, batches) for _ in 1:tmax] |> device
+    ys = [randn(rng, T, ny, batches) for _ in 1:tmax] |> device
+    x0 = init_states(model, batches) |> device
+
+    # Dummy loss function
+    function loss(model, x, us, ys)
+        m = is_diff ? model : REN(model)
+        J = 0
+        for t in 1:tmax
+            x, y = m(x, us[t])
+            J += Flux.mse(y, ys[t])
+        end
+        return J
+    end
+
+    # Run and time, running it once first to check it works
+    print("Forwards: ")
+    l = loss(model, x0, us, ys)
+    do_time && (@btime $loss($model, $x0, $us, $ys))
+
+    print("Reverse:  ")
+    g = gradient(loss, model, x0, us, ys)
+    do_time && (@btime $gradient($loss, $model, $x0, $us, $ys))
+
+    return l, g
+end
+
+# Test all types and combinations
+γ = 10
+ν = 10
+
+nu, nx, nv, ny = 4, 5, 10, 4
+X = randn(rng, ny, ny)
+Y = randn(rng, nu, nu)
+S = randn(rng, nu, ny)
+
+Q = -X'*X
+R = S * (Q \ S') + Y'*Y
+
+function test_rens(device)
+
+    d = device === cpu ? "CPU" : "GPU"
+    println("\nTesting RENs on ", d, ":")
+    println("--------------------\n")
+
+    println("Contracting REN:\n")
+    test_ren_device(device, ContractingRENParams)
+    println("\nContracting DiffREN:\n")
+    test_ren_device(device, ContractingRENParams; is_diff=true)
+
+    println("\nPassive REN:\n")
+    test_ren_device(device, PassiveRENParams, ν)
+    println("\nPassive DiffREN:\n")
+    test_ren_device(device, PassiveRENParams, ν; is_diff=true)
+
+    println("\nLipschitz REN:\n")
+    test_ren_device(device, LipschitzRENParams, γ)
+    println("\nLipschitz DiffREN:\n")
+    test_ren_device(device, LipschitzRENParams, γ; is_diff=true)
+
+    println("\nGeneral REN:\n")
+    test_ren_device(device, GeneralRENParams, Q, S, R)
+    println("\nGeneral DiffREN:\n")
+    test_ren_device(device, GeneralRENParams, Q, S, R; is_diff=true)
+
+    return nothing
+end
+
+test_rens(cpu)
+test_rens(gpu)

examples/GPU/test_sandwich.jl

@@ -0,0 +1,64 @@
+# This file is a part of RobustNeuralNetworks.jl. License is MIT: https://github.com/acfr/RobustNeuralNetworks.jl/blob/main/LICENSE 
+
+cd(@__DIR__)
+using Pkg
+Pkg.activate("../")
+
+using BenchmarkTools
+using CUDA
+using Flux
+using Random
+using RobustNeuralNetworks
+
+rng = Xoshiro(42)
+
+function test_sandwich_device(device; batches=400, do_time=true, T=Float32)
+
+    # Model parameters
+    nu = 2
+    nh = [10, 5]
+    ny = 4
+    γ = 10
+    nl = tanh
+
+    # Build model
+    model = Flux.Chain(
+        (x) -> (√γ * x),
+        SandwichFC(nu => nh[1], nl; T, rng),
+        SandwichFC(nh[1] => nh[2], nl; T, rng),
+        (x) -> (√γ * x),
+        SandwichFC(nh[2] => ny; output_layer=true, T, rng),
+    ) |> device
+
+    # Create dummy data
+    us = randn(rng, T, nu, batches) |> device
+    ys = randn(rng, T, ny, batches) |> device
+
+    # Dummy loss function
+    loss(model, u, y) = Flux.mse(model(u), y)
+
+    # Run and time, running it once to check it works
+    print("Forwards: ")
+    l = loss(model, us, ys)
+    do_time && (@btime $loss($model, $us, $ys))
+
+    print("Reverse:  ")
+    g = gradient(loss, model, us, ys)
+    do_time && (@btime $gradient($loss, $model, $us, $ys))
+
+    return l, g
+end
+
+function test_sandwich(device)
+
+    d = device === cpu ? "CPU" : "GPU"
+    println("\nTesting Sandwich on ", d, ":")
+    println("--------------------\n")
+
+    test_sandwich_device(device)
+
+    return nothing
+end
+
+test_sandwich(cpu)
+test_sandwich(gpu)

examples/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ControlSystemsBase = "aaaaaaaa-a6ca-5380-bf3e-84a91bcd477e"
@@ -17,3 +18,4 @@ Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 RobustNeuralNetworks = "a1f18e6b-8af1-433f-a85d-2e1ee636a2b8"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"