-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #118 from acfr/feature/gpu-support
Basic GPU Support
- Loading branch information
Showing
29 changed files
with
610 additions
and
324 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Binary file not shown.
Binary file not shown.
Oops, something went wrong.