From 5aeea6238c96f662ee578606c09507166a5c6aaa Mon Sep 17 00:00:00 2001 From: Mohamed Tarek Date: Sat, 14 Oct 2023 00:39:13 +1100 Subject: [PATCH 1/6] make CUDA an optional dependency --- Project.toml | 9 +++++++-- ext/MetalheadCUDAExt.jl | 21 +++++++++++++++++++++ src/layers/Layers.jl | 1 - src/layers/drop.jl | 7 ------- test/Project.toml | 1 + 5 files changed, 29 insertions(+), 10 deletions(-) create mode 100644 ext/MetalheadCUDAExt.jl diff --git a/Project.toml b/Project.toml index 6c4671f96..06e2acec4 100644 --- a/Project.toml +++ b/Project.toml @@ -1,11 +1,10 @@ name = "Metalhead" uuid = "dbeba491-748d-5e0e-a39e-b530a07fa0cc" -version = "0.8.3" +version = "0.9.0" [deps] Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0" -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" @@ -17,6 +16,12 @@ PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +[weakdeps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + +[extensions] +MetalheadCUDAExt = "CUDA" + [compat] BSON = "0.3.2" CUDA = "4, 5" diff --git a/ext/MetalheadCUDAExt.jl b/ext/MetalheadCUDAExt.jl new file mode 100644 index 000000000..5bad3914a --- /dev/null +++ b/ext/MetalheadCUDAExt.jl @@ -0,0 +1,21 @@ +module MetalheadCUDAExt + +if isdefined(Base, :get_extension) + using Metalhead + using Metalhead: Metalhead, _dropblock_mask + using CUDA: CUDA, CuArray +else + using ..Metalhead + using ..Metalhead: Metalhead, _dropblock_mask + using ..CUDA: CUDA, CuArray +end + +## bs is `clipped_block_size` +# Dispatch for GPU +Metalhead.dropblock_mask(rng::CUDA.RNG, x::CuArray, gamma, bs) = _dropblock_mask(rng, x, gamma, bs) +function Metalhead.dropblock_mask(rng, x::CuArray, gamma, bs) + throw(ArgumentError("x isa CuArray, but rng isa $(typeof(rng)). dropblock only supports + CUDA.RNG for CuArrays.")) +end + +end diff --git a/src/layers/Layers.jl b/src/layers/Layers.jl index 7c5dc0851..aaaa20228 100644 --- a/src/layers/Layers.jl +++ b/src/layers/Layers.jl @@ -2,7 +2,6 @@ module Layers using Flux using Flux: default_rng_value -using CUDA using NNlib using Functors using ChainRulesCore diff --git a/src/layers/drop.jl b/src/layers/drop.jl index 15f8e7533..8593f9301 100644 --- a/src/layers/drop.jl +++ b/src/layers/drop.jl @@ -43,13 +43,6 @@ function dropblock(rng::AbstractRNG, x::AbstractArray{T, 4}, drop_block_prob, return x .* block_mask .* normalize_scale end -## bs is `clipped_block_size` -# Dispatch for GPU -dropblock_mask(rng::CUDA.RNG, x::CuArray, gamma, bs) = _dropblock_mask(rng, x, gamma, bs) -function dropblock_mask(rng, x::CuArray, gamma, bs) - throw(ArgumentError("x isa CuArray, but rng isa $(typeof(rng)). dropblock only supports - CUDA.RNG for CuArrays.")) -end # Dispatch for CPU dropblock_mask(rng, x, gamma, bs) = _dropblock_mask(rng, x, gamma, bs) diff --git a/test/Project.toml b/test/Project.toml index e65c612a4..b121720e7 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0" ReTest = "e0db7c4e-2690-44b9-bad6-7687da720f89" From e1e8efb8021bbf3d75d544797385dae27ad1b3e3 Mon Sep 17 00:00:00 2001 From: Mohamed Tarek Date: Sat, 14 Oct 2023 00:43:49 +1100 Subject: [PATCH 2/6] simplify --- ext/MetalheadCUDAExt.jl | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/ext/MetalheadCUDAExt.jl b/ext/MetalheadCUDAExt.jl index 5bad3914a..84b2c9ed9 100644 --- a/ext/MetalheadCUDAExt.jl +++ b/ext/MetalheadCUDAExt.jl @@ -1,14 +1,8 @@ module MetalheadCUDAExt -if isdefined(Base, :get_extension) - using Metalhead - using Metalhead: Metalhead, _dropblock_mask - using CUDA: CUDA, CuArray -else - using ..Metalhead - using ..Metalhead: Metalhead, _dropblock_mask - using ..CUDA: CUDA, CuArray -end +using Metalhead +using Metalhead: Metalhead, _dropblock_mask +using CUDA: CUDA, CuArray ## bs is `clipped_block_size` # Dispatch for GPU From d6308a0d6f62539b77e3689322733427c25824d3 Mon Sep 17 00:00:00 2001 From: Mohamed Tarek Date: Sat, 14 Oct 2023 02:01:29 +1100 Subject: [PATCH 3/6] make CUDA both a dep and an optional dep --- Project.toml | 1 + src/Metalhead.jl | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/Project.toml b/Project.toml index 06e2acec4..bf730f532 100644 --- a/Project.toml +++ b/Project.toml @@ -6,6 +6,7 @@ version = "0.9.0" Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" diff --git a/src/Metalhead.jl b/src/Metalhead.jl index d7db8d4c3..ced4ca446 100644 --- a/src/Metalhead.jl +++ b/src/Metalhead.jl @@ -93,4 +93,8 @@ for T in (:AlexNet, :VGG, :SqueezeNet, :ResNet, :WideResNet, :ResNeXt, @eval Base.show(io::IO, ::MIME"text/plain", model::$T) = _maybe_big_show(io, model) end +if !isdefined(Base, :get_extension) + include("../ext/MetalheadCUDAExt.jl") +end + end # module From c9de3c060f9dc9060b9f31feac7498053ae84d5c Mon Sep 17 00:00:00 2001 From: Mohamed Tarek Date: Sat, 14 Oct 2023 02:03:03 +1100 Subject: [PATCH 4/6] patch version bump --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index bf730f532..a24fe91c7 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Metalhead" uuid = "dbeba491-748d-5e0e-a39e-b530a07fa0cc" -version = "0.9.0" +version = "0.8.4" [deps] Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" From 998af04ec6cd1e1daaa0deea390e746db8779f26 Mon Sep 17 00:00:00 2001 From: Mohamed Tarek Date: Sat, 14 Oct 2023 02:08:31 +1100 Subject: [PATCH 5/6] fix loading packages in ext --- ext/MetalheadCUDAExt.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/ext/MetalheadCUDAExt.jl b/ext/MetalheadCUDAExt.jl index 84b2c9ed9..35d268da9 100644 --- a/ext/MetalheadCUDAExt.jl +++ b/ext/MetalheadCUDAExt.jl @@ -1,7 +1,10 @@ module MetalheadCUDAExt -using Metalhead -using Metalhead: Metalhead, _dropblock_mask +if isdefined(Base, :get_extension) + using Metalhead: Metalhead, _dropblock_mask +else + using ..Metalhead: Metalhead, _dropblock_mask +end using CUDA: CUDA, CuArray ## bs is `clipped_block_size` From dc4fb0a285d20d646cd13f195c625e8ef8af5b9d Mon Sep 17 00:00:00 2001 From: Mohamed Tarek Date: Sat, 14 Oct 2023 03:02:10 +1100 Subject: [PATCH 6/6] Metalhead. -> Metalhead.Layers. --- ext/MetalheadCUDAExt.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ext/MetalheadCUDAExt.jl b/ext/MetalheadCUDAExt.jl index 35d268da9..72fa126d9 100644 --- a/ext/MetalheadCUDAExt.jl +++ b/ext/MetalheadCUDAExt.jl @@ -1,16 +1,16 @@ module MetalheadCUDAExt if isdefined(Base, :get_extension) - using Metalhead: Metalhead, _dropblock_mask + using Metalhead: Metalhead else - using ..Metalhead: Metalhead, _dropblock_mask + using ..Metalhead: Metalhead end using CUDA: CUDA, CuArray ## bs is `clipped_block_size` # Dispatch for GPU -Metalhead.dropblock_mask(rng::CUDA.RNG, x::CuArray, gamma, bs) = _dropblock_mask(rng, x, gamma, bs) -function Metalhead.dropblock_mask(rng, x::CuArray, gamma, bs) +Metalhead.Layers.dropblock_mask(rng::CUDA.RNG, x::CuArray, gamma, bs) = Metalhead.Layers._dropblock_mask(rng, x, gamma, bs) +function Metalhead.Layers.dropblock_mask(rng, x::CuArray, gamma, bs) throw(ArgumentError("x isa CuArray, but rng isa $(typeof(rng)). dropblock only supports CUDA.RNG for CuArrays.")) end