add some grad tests on GPU
This commit is contained in:
Dhairya Gandhi
parent
59bb0d81b0
commit
717ad9328d
|
|
@ -52,9 +52,9 @@ version = "4.0.4"
|
|||
|
||||
[[CUDAnative]]
|
||||
deps = ["Adapt", "CEnum", "CUDAapi", "CUDAdrv", "DataStructures", "InteractiveUtils", "LLVM", "Libdl", "Printf", "TimerOutputs"]
|
||||
git-tree-sha1 = "93f6c917ab2a9b5bb54f8f738f4ec1a6693cb716"
|
||||
git-tree-sha1 = "a67b38619d1fa131027bac1c4a81f0012254d1fd"
|
||||
uuid = "be33ccc6-a3ff-5ff2-a52e-74243cff1e17"
|
||||
version = "2.5.5"
|
||||
version = "2.6.0"
|
||||
|
||||
[[CodecZlib]]
|
||||
deps = ["BinaryProvider", "Libdl", "TranscodingStreams"]
|
||||
|
|
@ -94,9 +94,11 @@ version = "1.3.0"
|
|||
|
||||
[[CuArrays]]
|
||||
deps = ["AbstractFFTs", "Adapt", "CEnum", "CUDAapi", "CUDAdrv", "CUDAnative", "DataStructures", "GPUArrays", "Libdl", "LinearAlgebra", "MacroTools", "NNlib", "Printf", "Random", "Requires", "SparseArrays", "TimerOutputs"]
|
||||
git-tree-sha1 = "4757376a85ffb27d4c4f6cdf9635261e6c3a5fec"
|
||||
git-tree-sha1 = "dc6e4f2bf3d84861ca00face3f0209ab71a0bc00"
|
||||
repo-rev = "master"
|
||||
repo-url = "https://github.com/JuliaGPU/CuArrays.jl.git"
|
||||
uuid = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
|
||||
version = "1.4.7"
|
||||
version = "1.5.0"
|
||||
|
||||
[[DataAPI]]
|
||||
git-tree-sha1 = "674b67f344687a88310213ddfa8a2b3c76cc4252"
|
||||
|
|
@ -204,10 +206,10 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
|
|||
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
|
||||
|
||||
[[MacroTools]]
|
||||
deps = ["Compat", "DataStructures", "Test"]
|
||||
git-tree-sha1 = "82921f0e3bde6aebb8e524efc20f4042373c0c06"
|
||||
deps = ["DataStructures", "Markdown", "Random"]
|
||||
git-tree-sha1 = "e2fc7a55bb2224e203bbd8b59f72b91323233458"
|
||||
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
|
||||
version = "0.5.2"
|
||||
version = "0.5.3"
|
||||
|
||||
[[Markdown]]
|
||||
deps = ["Base64"]
|
||||
|
|
@ -229,8 +231,10 @@ version = "0.4.3"
|
|||
uuid = "a63ad114-7e13-5084-954f-fe012c677804"
|
||||
|
||||
[[NNlib]]
|
||||
deps = ["Libdl", "LinearAlgebra", "Requires", "Statistics", "TimerOutputs"]
|
||||
git-tree-sha1 = "0c667371391fc6bb31f7f12f96a56a17098b3de8"
|
||||
deps = ["BinaryProvider", "Libdl", "LinearAlgebra", "Requires", "Statistics"]
|
||||
git-tree-sha1 = "fe5de1375620c3b5e04bdfc3dee5199cbefb2c0c"
|
||||
repo-rev = "master"
|
||||
repo-url = "https://github.com/FluxML/NNlib.jl.git"
|
||||
uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
|
||||
version = "0.6.0"
|
||||
|
||||
|
|
|
|||
|
|
@ -61,6 +61,7 @@ if CuArrays.has_cudnn()
|
|||
@info "Testing Flux/CUDNN"
|
||||
include("cudnn.jl")
|
||||
include("curnn.jl")
|
||||
include("layers.jl")
|
||||
else
|
||||
@warn "CUDNN unavailable, not testing GPU DNN support"
|
||||
end
|
||||
|
|
|
|||
|
|
@ -0,0 +1,86 @@
|
|||
# Test layers and data/model movements on and off the GPU
|
||||
# Add tests for layers and their gradients on the GPU
|
||||
# Most of the forward passes should be fine being applied
|
||||
# to bitstype objects, but this gives higher coverage for our use-cases
|
||||
# Check that getting the gradients does not throw
|
||||
|
||||
# generic movement tests
|
||||
@test_broken gradient(x -> sum(gpu(x)), rand(3,3)) isa Tuple
|
||||
@test gradient(x -> sum(cpu(x)), gpu(rand(3,3))) isa Tuple
|
||||
|
||||
function gradtest(layers, args...; name = "Conv", xs = rand(Float32, 28, 28, 1, 1))
|
||||
@testset "$name GPU grad tests" begin
|
||||
for layer in layers
|
||||
@testset "$layer GPU grad test" begin
|
||||
l = gpu(layer(args...))
|
||||
xs = gpu(xs)
|
||||
if l isa DepthwiseConv
|
||||
@test_broken gradient(Flux.params(l)) do
|
||||
sum(l(xs))
|
||||
end isa Flux.Zygote.Grads
|
||||
else
|
||||
@test gradient(Flux.params(l)) do
|
||||
sum(l(xs))
|
||||
end isa Flux.Zygote.Grads
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
# Repeats from Conv, CrossCor
|
||||
|
||||
conv_layers = [Conv, ConvTranspose, CrossCor, DepthwiseConv]
|
||||
gradtest(conv_layers, (2,2), 1=>3, name = "Conv")
|
||||
|
||||
pooling_layers = [MaxPool, MeanPool]
|
||||
gradtest(pooling_layers, (2,2), name = "Pooling")
|
||||
|
||||
dropout_layers = [Dropout, AlphaDropout]
|
||||
gradtest(dropout_layers, 0.5f0, name = "Dropout")
|
||||
|
||||
norm_layers = [LayerNorm, BatchNorm]
|
||||
gradtest(norm_layers, 1, name = "Normalising", xs = rand(Float32, 28,28,3,1))
|
||||
|
||||
instancenorm = [InstanceNorm]
|
||||
gradtest(instancenorm, 1, name = "InstanceNorm")
|
||||
|
||||
groupnorm = [GroupNorm]
|
||||
gradtest(groupnorm, 3, 1, name = "GroupNorm", xs = rand(Float32, 28,28,3,1))
|
||||
|
||||
const stateless_layers = [Flux.mse,
|
||||
Flux.crossentropy,
|
||||
Flux.logitcrossentropy,]
|
||||
Flux.normalise]
|
||||
|
||||
const stateless_layers_broadcasted = [Flux.binarycrossentropy,
|
||||
Flux.logitbinarycrossentropy]
|
||||
|
||||
function stateless_gradtest(f, args...)
|
||||
@test gradient((args...) -> sum(f(args...)), args...)[1] isa CuArray
|
||||
end
|
||||
|
||||
function stateless_gradtest_broadcasted(f, args...)
|
||||
if f == Flux.binarycrossentropy
|
||||
@test_broken gradient((args...) -> sum(f.(args...)), args...)[1] isa CuArray
|
||||
else
|
||||
@test gradient((args...) -> sum(f.(args...)), args...)[1] isa CuArray
|
||||
end
|
||||
end
|
||||
|
||||
@testset "Stateless GPU grad tests" begin
|
||||
x = gpu(rand(3,3))
|
||||
y = gpu(rand(3,3))
|
||||
|
||||
for layer in stateless_layers
|
||||
if layer == Flux.normalise
|
||||
stateless_gradtest(layer, x)
|
||||
else
|
||||
stateless_gradtest(layer, x, y)
|
||||
end
|
||||
end
|
||||
|
||||
for layer in stateless_layers_broadcasted
|
||||
stateless_gradtest_broadcasted(layer, x, y)
|
||||
end
|
||||
end
|
||||
Loading…
Reference in New Issue