Skip to content

Commit 8772a02

Browse files
CarloLucibelloCarloLucibello
committed
gpu tests
1 parent 9029377 commit 8772a02

File tree

4 files changed

+269
-9
lines changed

4 files changed

+269
-9
lines changed

‎GraphNeuralNetworks/Project.toml‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,28 +33,26 @@ LinearAlgebra = "1"
3333
MLUtils = "0.4"
3434
MacroTools = "0.5"
3535
NNlib = "0.9"
36+
Pkg = "1"
3637
Random = "1"
3738
Reexport = "1"
3839
Statistics = "1"
3940
TestItemRunner = "1.0.5"
40-
cuDNN = "1"
4141
julia = "1.10"
4242

4343
[extras]
4444
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
45-
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
4645
ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
4746
DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
4847
FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
4948
Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
5049
InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
5150
MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
51+
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
5252
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
5353
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
5454
TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
5555
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
56-
cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
5756

5857
[targets]
59-
test = ["Test", "TestItemRunner", "MLDatasets", "Adapt", "DataFrames", "InlineStrings",
60-
"SparseArrays", "Graphs", "Zygote", "FiniteDifferences", "ChainRulesTestUtils", "CUDA", "cuDNN"]
58+
test = ["Test", "TestItemRunner", "Pkg", "MLDatasets", "Adapt", "DataFrames", "InlineStrings", "SparseArrays", "Graphs", "Zygote", "FiniteDifferences", "ChainRulesTestUtils"]

‎GraphNeuralNetworks/test/layers/conv.jl‎

Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,17 @@ end
6565
end
6666
end
6767

68+
69+
@testitem "GCNConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
70+
using .TestModule
71+
l = GCNConv(D_IN => D_OUT)
72+
for g in TEST_GRAPHS
73+
g.graph isa AbstractSparseMatrix && continue
74+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
75+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
76+
end
77+
end
78+
6879
@testitem "ChebConv" setup=[TolSnippet, TestModule] begin
6980
using .TestModule
7081
k = 2
@@ -84,6 +95,18 @@ end
8495
end
8596
end
8697

98+
99+
@testitem "ChebConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
100+
using .TestModule
101+
k = 2
102+
l = ChebConv(D_IN => D_OUT, k)
103+
for g in TEST_GRAPHS
104+
g.graph isa AbstractSparseMatrix && continue
105+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
106+
test_gradients(l, g, g.x, rtol = RTOL_LOW, test_gpu = true, compare_finite_diff = false)
107+
end
108+
end
109+
87110
@testitem "GraphConv" setup=[TolSnippet, TestModule] begin
88111
using .TestModule
89112
l = GraphConv(D_IN => D_OUT)
@@ -104,6 +127,18 @@ end
104127
end
105128
end
106129

130+
131+
@testitem "GraphConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
132+
using .TestModule
133+
l = GraphConv(D_IN => D_OUT)
134+
for g in TEST_GRAPHS
135+
g.graph isa AbstractSparseMatrix && continue
136+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
137+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
138+
end
139+
end
140+
141+
107142
@testitem "GATConv" setup=[TolSnippet, TestModule] begin
108143
using .TestModule
109144
for heads in (1, 2), concat in (true, false)
@@ -132,6 +167,18 @@ end
132167
end
133168
end
134169

170+
@testitem "GATConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
171+
using .TestModule
172+
for heads in (1, 2), concat in (true, false)
173+
l = GATConv(D_IN => D_OUT; heads, concat, dropout=0)
174+
for g in TEST_GRAPHS
175+
g.graph isa AbstractSparseMatrix && continue
176+
@test size(l(g, g.x)) == (concat ? heads * D_OUT : D_OUT, g.num_nodes)
177+
test_gradients(l, g, g.x, rtol = RTOL_LOW, test_gpu = true, compare_finite_diff = false)
178+
end
179+
end
180+
end
181+
135182
@testitem "GATv2Conv" setup=[TolSnippet, TestModule] begin
136183
using .TestModule
137184
for heads in (1, 2), concat in (true, false)
@@ -160,6 +207,18 @@ end
160207
end
161208
end
162209

210+
@testitem "GATv2Conv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
211+
using .TestModule
212+
for heads in (1, 2), concat in (true, false)
213+
l = GATv2Conv(D_IN => D_OUT, tanh; heads, concat, dropout=0)
214+
for g in TEST_GRAPHS
215+
g.graph isa AbstractSparseMatrix && continue
216+
@test size(l(g, g.x)) == (concat ? heads * D_OUT : D_OUT, g.num_nodes)
217+
test_gradients(l, g, g.x, rtol = RTOL_LOW, atol=ATOL_LOW, test_gpu = true, compare_finite_diff = false)
218+
end
219+
end
220+
end
221+
163222
@testitem "GatedGraphConv" setup=[TolSnippet, TestModule] begin
164223
using .TestModule
165224
num_layers = 3
@@ -172,6 +231,18 @@ end
172231
end
173232
end
174233

234+
235+
@testitem "GatedGraphConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
236+
using .TestModule
237+
num_layers = 3
238+
l = GatedGraphConv(D_OUT, num_layers)
239+
for g in TEST_GRAPHS
240+
g.graph isa AbstractSparseMatrix && continue
241+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
242+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
243+
end
244+
end
245+
175246
@testitem "EdgeConv" setup=[TolSnippet, TestModule] begin
176247
using .TestModule
177248
l = EdgeConv(Dense(2 * D_IN, D_OUT), aggr = +)
@@ -194,6 +265,30 @@ end
194265
@test !in(:eps, Flux.trainable(l))
195266
end
196267

268+
@testitem "GINConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
269+
using .TestModule
270+
nn = Dense(D_IN, D_OUT)
271+
l = GINConv(nn, 0.01, aggr = mean)
272+
for g in TEST_GRAPHS
273+
g.graph isa AbstractSparseMatrix && continue
274+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
275+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
276+
end
277+
end
278+
279+
@testitem "NNConv" setup=[TolSnippet, TestModule] begin
280+
using .TestModule
281+
edim = 10
282+
nn = Dense(edim, D_OUT * D_IN)
283+
284+
l = NNConv(D_IN => D_OUT, nn, tanh, bias = true, aggr = +)
285+
for g in TEST_GRAPHS
286+
g = GNNGraph(g, edata = rand(Float32, edim, g.num_edges))
287+
@test size(l(g, g.x, g.e)) == (D_OUT, g.num_nodes)
288+
test_gradients(l, g, g.x, g.e, rtol = RTOL_HIGH)
289+
end
290+
end
291+
197292
@testitem "NNConv" setup=[TolSnippet, TestModule] begin
198293
using .TestModule
199294
edim = 10
@@ -219,6 +314,16 @@ end
219314
end
220315
end
221316

317+
@testitem "SAGEConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
318+
using .TestModule
319+
l = SAGEConv(D_IN => D_OUT)
320+
for g in TEST_GRAPHS
321+
g.graph isa AbstractSparseMatrix && continue
322+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
323+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
324+
end
325+
end
326+
222327
@testitem "ResGatedGraphConv" setup=[TolSnippet, TestModule] begin
223328
using .TestModule
224329
l = ResGatedGraphConv(D_IN => D_OUT, tanh, bias = true)
@@ -228,6 +333,16 @@ end
228333
end
229334
end
230335

336+
@testitem "ResGatedGraphConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
337+
using .TestModule
338+
l = ResGatedGraphConv(D_IN => D_OUT, tanh, bias = true)
339+
for g in TEST_GRAPHS
340+
g.graph isa AbstractSparseMatrix && continue
341+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
342+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
343+
end
344+
end
345+
231346
@testitem "CGConv" setup=[TolSnippet, TestModule] begin
232347
using .TestModule
233348

@@ -246,6 +361,17 @@ end
246361
@test l1(g1, g1.ndata.x, nothing) == l1(g1).ndata.x
247362
end
248363

364+
@testitem "CGConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
365+
using .TestModule
366+
edim = 10
367+
l = CGConv((D_IN, edim) => D_OUT, tanh, residual = false, bias = true)
368+
for g in TEST_GRAPHS
369+
g.graph isa AbstractSparseMatrix && continue
370+
@test size(l(g, g.x, g.e)) == (D_OUT, g.num_nodes)
371+
test_gradients(l, g, g.x, g.e, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
372+
end
373+
end
374+
249375
@testitem "AGNNConv" setup=[TolSnippet, TestModule] begin
250376
using .TestModule
251377
l = AGNNConv(trainable=false, add_self_loops=false)
@@ -265,6 +391,16 @@ end
265391
end
266392
end
267393

394+
@testitem "AGNNConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
395+
using .TestModule
396+
l = AGNNConv(trainable=false, add_self_loops=false)
397+
for g in TEST_GRAPHS
398+
g.graph isa AbstractSparseMatrix && continue
399+
@test size(l(g, g.x)) == (D_IN, g.num_nodes)
400+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
401+
end
402+
end
403+
268404
@testitem "MEGNetConv" setup=[TolSnippet, TestModule] begin
269405
using .TestModule
270406
l = MEGNetConv(D_IN => D_OUT, aggr = +)
@@ -281,6 +417,22 @@ end
281417
end
282418
end
283419

420+
@testitem "MEGNetConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
421+
using .TestModule
422+
l = MEGNetConv(D_IN => D_OUT, aggr = +)
423+
for g in TEST_GRAPHS
424+
g.graph isa AbstractSparseMatrix && continue
425+
y = l(g, g.x, g.e)
426+
@test size(y[1]) == (D_OUT, g.num_nodes)
427+
@test size(y[2]) == (D_OUT, g.num_edges)
428+
function loss(l, g, x, e)
429+
y = l(g, x, e)
430+
return mean(y[1]) + sum(y[2])
431+
end
432+
test_gradients(l, g, g.x, g.e, rtol = RTOL_LOW; loss, test_gpu = true, compare_finite_diff = false)
433+
end
434+
end
435+
284436
@testitem "GMMConv" setup=[TolSnippet, TestModule] begin
285437
using .TestModule
286438
ein_channel = 10
@@ -293,6 +445,18 @@ end
293445
end
294446
end
295447

448+
@testitem "GMMConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
449+
using .TestModule
450+
ein_channel = 10
451+
K = 5
452+
l = GMMConv((D_IN, ein_channel) => D_OUT, K = K)
453+
for g in TEST_GRAPHS
454+
g.graph isa AbstractSparseMatrix && continue
455+
y = l(g, g.x, g.e)
456+
test_gradients(l, g, g.x, g.e, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
457+
end
458+
end
459+
296460
@testitem "SGConv" setup=[TolSnippet, TestModule] begin
297461
using .TestModule
298462
K = [1, 2, 3] # for different number of hops
@@ -311,6 +475,17 @@ end
311475
end
312476
end
313477

478+
@testitem "SGConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
479+
using .TestModule
480+
k = 2
481+
l = SGConv(D_IN => D_OUT, k, add_self_loops = true)
482+
for g in TEST_GRAPHS
483+
g.graph isa AbstractSparseMatrix && continue
484+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
485+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
486+
end
487+
end
488+
314489
@testitem "TAGConv" setup=[TolSnippet, TestModule] begin
315490
using .TestModule
316491
K = [1, 2, 3]
@@ -329,9 +504,21 @@ end
329504
end
330505
end
331506

507+
@testitem "TAGConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
508+
using .TestModule
509+
k = 2
510+
l = TAGConv(D_IN => D_OUT, k, add_self_loops = true)
511+
for g in TEST_GRAPHS
512+
g.graph isa AbstractSparseMatrix && continue
513+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
514+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
515+
end
516+
end
517+
332518
@testitem "EGNNConv" setup=[TolSnippet, TestModule] begin
333519
using .TestModule
334520
#TODO test gradient
521+
#TODO test gpu
335522
@testset "EGNNConv $GRAPH_T" for GRAPH_T in GRAPH_TYPES
336523
hin = 5
337524
hout = 5
@@ -378,6 +565,22 @@ end
378565
end
379566
end
380567

568+
@testitem "TransformerConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
569+
using .TestModule
570+
ein = 2
571+
heads = 3
572+
573+
# used like in Shi et al., 2021
574+
l = TransformerConv((D_IN, ein) => D_IN; heads, gating = true,
575+
bias_qkv = true)
576+
for g in TEST_GRAPHS
577+
g.graph isa AbstractSparseMatrix && continue
578+
@test size(l(g, g.x, g.e)) == (D_IN * heads, g.num_nodes)
579+
test_gradients(l, g, g.x, g.e, rtol = RTOL_LOW, test_gpu = true, compare_finite_diff = false)
580+
end
581+
end
582+
583+
381584
@testitem "DConv" setup=[TolSnippet, TestModule] begin
382585
using .TestModule
383586
K = [1, 2, 3] # for different number of hops
@@ -389,3 +592,13 @@ end
389592
end
390593
end
391594
end
595+
596+
@testitem "DConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
597+
using .TestModule
598+
l = DConv(D_IN => D_OUT, 2)
599+
for g in TEST_GRAPHS
600+
g.graph isa AbstractSparseMatrix && continue
601+
@test size(l(g, g.x)) == (D_OUT, g.num_nodes)
602+
test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
603+
end
604+
end

‎GraphNeuralNetworks/test/runtests.jl‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,31 @@
11
using TestItemRunner
2+
using Pkg
23
# See https://www.julia-vscode.org/docs/stable/userguide/testitems/
34
# for how to run the tests within VS Code.
4-
@run_package_tests
5+
# See test_module.jl for the test infrasructure.
6+
7+
## Uncomment below to change the default test settings
8+
# ENV["GNN_TEST_CPU"] = "false"
9+
# ENV["GNN_TEST_CUDA"] = "true"
10+
# ENV["GNN_TEST_AMDGPU"] = "true"
11+
# ENV["GNN_TEST_Metal"] = "true"
12+
13+
if get(ENV, "GNN_TEST_CPU", "true") == "true"
14+
@run_package_tests filter = ti -> :gpu ∉ ti.tags
15+
end
16+
if get(ENV, "GNN_TEST_CUDA", "false") == "true"
17+
Pkg.add(["CUDA", "cuDNN"])
18+
using CUDA
19+
@run_package_tests filter = ti -> :gpu ∈ ti.tags
20+
end
21+
if get(ENV, "GNN_TEST_AMDGPU", "false") == "true"
22+
Pkg.add("AMDGPU")
23+
using AMDGPU
24+
@run_package_tests filter = ti -> :gpu ∈ ti.tags
25+
end
26+
if get(ENV, "GNN_TEST_Metal", "false") == "true"
27+
Pkg.add("Metal")
28+
using Metal
29+
@run_package_tests filter = ti -> :gpu ∈ ti.tags
30+
end
31+

0 commit comments

Comments
 (0)