Skip to content
This repository was archived by the owner on Sep 28, 2024. It is now read-only.

Commit 65e0749

Browse files
yuehhuayuehhua
committed
fix DeepONet for CUDA
1 parent cf8f4bd commit 65e0749

File tree

5 files changed

+88
-56
lines changed

5 files changed

+88
-56
lines changed

‎example/Burgers/src/Burgers_deeponet.jl‎

Lines changed: 21 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,34 @@
1-
function train_don()
2-
# if has_cuda()
3-
# @info "CUDA is on"
4-
# device = gpu
5-
# CUDA.allowscalar(false)
6-
# else
1+
function train_don(; n=300, cuda=true, learning_rate=0.001, epochs=400)
2+
if cuda && has_cuda()
3+
@info "Training on GPU"
4+
device = gpu
5+
else
6+
@info "Training on CPU"
77
device = cpu
8-
# end
8+
end
99

10-
x, y = get_data_don(n=300)
11-
xtrain = x[1:280, :]' |> device
12-
xval = x[end-19:end, :]' |> device
10+
x, y = get_data_don(n=n)
11+
12+
xtrain = x[1:280, :]'
13+
ytrain = y[1:280, :]
1314

14-
ytrain = y[1:280, :] |> device
15+
xval = x[end-19:end, :]' |> device
1516
yval = y[end-19:end, :] |> device
1617

17-
grid = collect(range(0, 1, length=1024))' |> device
18+
grid = collect(range(0, 1, length=1024)') |> device
1819

19-
learning_rate = 0.001
2020
opt = ADAM(learning_rate)
2121

22-
m = DeepONet((1024,1024,1024),(1,1024,1024),gelu,gelu)
23-
loss(xtrain,ytrain,sensor) = Flux.Losses.mse(m(xtrain,sensor),ytrain)
24-
evalcb() = @show(loss(xval,yval,grid))
22+
m = DeepONet((1024,1024,1024), (1,1024,1024), gelu, gelu) |> device
23+
24+
loss(X, y, sensor) = Flux.Losses.mse(m(X, sensor), y)
25+
evalcb() = @show(loss(xval, yval, grid))
2526

26-
Flux.@epochs 400 Flux.train!(loss, params(m), [(xtrain,ytrain,grid)], opt, cb = evalcb)
27-
ỹ = m(xval, grid)
27+
data = [(xtrain, ytrain, grid)] |> device
28+
Flux.@epochs epochs Flux.train!(loss, params(m), data, opt, cb=evalcb)
29+
ỹ = m(xval |> device, grid |> device)
2830

29-
diffvec = vec(abs.((yval .- ỹ)))
31+
diffvec = vec(abs.(cpu(yval) .- cpu(ỹ)))
3032
mean_diff = sum(diffvec)/length(diffvec)
3133
return mean_diff
3234
end

‎src/DeepONet.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ function (a::DeepONet)(x::AbstractArray, y::AbstractVecOrMat)
116116
However, we perform the transformations by the NNs always in the first dim
117117
so we need to adjust (i.e. transpose) one of the inputs,
118118
which we do on the branch input here =#
119-
return Array(branch(x)') * trunk(y)
119+
return branch(x)' * trunk(y)
120120
end
121121

122122
# Sensors stay the same and shouldn't be batched

‎test/cuda.jl‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
@testset "CUDA" begin
2+
@testset "DeepONet" begin
3+
batch_size = 2
4+
a = [0.83541104, 0.83479851, 0.83404712, 0.83315711, 0.83212979, 0.83096755,
5+
0.82967374, 0.82825263, 0.82670928, 0.82504949, 0.82327962, 0.82140651,
6+
0.81943734, 0.81737952, 0.8152405, 0.81302771]
7+
a = repeat(a, outer=(1, batch_size)) |> gpu
8+
sensors = collect(range(0, 1, length=16)')
9+
sensors = repeat(sensors, outer=(batch_size, 1)) |> gpu
10+
model = DeepONet((16, 22, 30), (2, 16, 24, 30), σ, tanh;
11+
init_branch=Flux.glorot_normal, bias_trunk=false) |> gpu
12+
y = model(a, sensors)
13+
@test size(y) == (batch_size, 16)
14+
15+
mgrad = Flux.Zygote.gradient(() -> sum(model(a, sensors)), Flux.params(model))
16+
@test length(mgrad.grads) == 9
17+
end
18+
end

‎test/deeponet.jl‎

Lines changed: 25 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,34 @@
1-
using Test, Flux
2-
31
@testset "DeepONet" begin
4-
@testset "dimensions" begin
5-
# Test the proper construction
2+
@testset "proper construction" begin
3+
deeponet = DeepONet((32,64,72), (24,48,72), σ, tanh)
64
# Branch net
7-
@test size(DeepONet((32,64,72), (24,48,72), σ, tanh).branch_net.layers[end].weight) == (72,64)
8-
@test size(DeepONet((32,64,72), (24,48,72), σ, tanh).branch_net.layers[end].bias) == (72,)
5+
@test size(deeponet.branch_net.layers[end].weight) == (72,64)
6+
@test size(deeponet.branch_net.layers[end].bias) == (72,)
97
# Trunk net
10-
@test size(DeepONet((32,64,72), (24,48,72), σ, tanh).trunk_net.layers[end].weight) == (72,48)
11-
@test size(DeepONet((32,64,72), (24,48,72), σ, tanh).trunk_net.layers[end].bias) == (72,)
8+
@test size(deeponet.trunk_net.layers[end].weight) == (72,48)
9+
@test size(deeponet.trunk_net.layers[end].bias) == (72,)
1210
end
1311

1412
# Accept only Int as architecture parameters
1513
@test_throws MethodError DeepONet((32.5,64,72), (24,48,72), σ, tanh)
1614
@test_throws MethodError DeepONet((32,64,72), (24.1,48,72))
17-
end
18-
19-
#Just the first 16 datapoints from the Burgers' equation dataset
20-
a = [0.83541104, 0.83479851, 0.83404712, 0.83315711, 0.83212979, 0.83096755, 0.82967374, 0.82825263, 0.82670928, 0.82504949, 0.82327962, 0.82140651, 0.81943734, 0.81737952, 0.8152405, 0.81302771]
21-
sensors = collect(range(0, 1, length=16))'
22-
23-
model = DeepONet((16, 22, 30), (1, 16, 24, 30), σ, tanh; init_branch=Flux.glorot_normal, bias_trunk=false)
24-
25-
model(a,sensors)
2615

27-
#forward pass
28-
@test size(model(a, sensors)) == (1, 16)
29-
30-
mgrad = Flux.Zygote.gradient((x,p)->sum(model(x,p)),a,sensors)
31-
32-
#gradients
33-
@test !iszero(Flux.Zygote.gradient((x,p)->sum(model(x,p)),a,sensors)[1])
34-
@test !iszero(Flux.Zygote.gradient((x,p)->sum(model(x,p)),a,sensors)[2])
35-
36-
#Output size of branch and trunk subnets should be same
37-
branch = Chain(Dense(16, 22), Dense(22, 30))
38-
trunk = Chain(Dense(1, 16), Dense(16, 24), Dense(24, 32))
39-
m = DeepONet(branch, trunk)
40-
@test_throws AssertionError DeepONet((32,64,70), (24,48,72), σ, tanh)
41-
@test_throws DimensionMismatch m(a, sensors)
16+
# Just the first 16 datapoints from the Burgers' equation dataset
17+
a = [0.83541104, 0.83479851, 0.83404712, 0.83315711, 0.83212979, 0.83096755,
18+
0.82967374, 0.82825263, 0.82670928, 0.82504949, 0.82327962, 0.82140651,
19+
0.81943734, 0.81737952, 0.8152405, 0.81302771]
20+
sensors = collect(range(0, 1, length=16)')
21+
model = DeepONet((16, 22, 30), (1, 16, 24, 30), σ, tanh; init_branch=Flux.glorot_normal, bias_trunk=false)
22+
y = model(a, sensors)
23+
@test size(y) == (1, 16)
24+
25+
mgrad = Flux.Zygote.gradient(() -> sum(model(a, sensors)), Flux.params(model))
26+
@test length(mgrad.grads) == 7
27+
28+
# Output size of branch and trunk subnets should be same
29+
branch = Chain(Dense(16, 22), Dense(22, 30))
30+
trunk = Chain(Dense(1, 16), Dense(16, 24), Dense(24, 32))
31+
m = DeepONet(branch, trunk)
32+
@test_throws AssertionError DeepONet((32,64,70), (24,48,72), σ, tanh)
33+
@test_throws DimensionMismatch m(a, sensors)
34+
end

‎test/runtests.jl‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,31 @@
11
using NeuralOperators
22
using Test
33
using Flux
4+
using CUDA
5+
6+
CUDA.allowscalar(false)
7+
8+
cuda_tests = [
9+
"cuda",
10+
]
11+
12+
tests = [
13+
"Transform/Transform",
14+
"operator_kernel",
15+
"model",
16+
"deeponet",
17+
]
18+
19+
if CUDA.functional()
20+
append!(tests, cuda_tests)
21+
else
22+
@warn "CUDA unavailable, not testing GPU support"
23+
end
424

525
@testset "NeuralOperators.jl" begin
6-
include("Transform/Transform.jl")
7-
include("operator_kernel.jl")
8-
include("model.jl")
9-
include("deeponet.jl")
26+
for t in tests
27+
include("$(t).jl")
28+
end
1029
end
1130

1231
#=

0 commit comments

Comments
 (0)