TastyPancakes
/
Flux.jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

test gradients are allocated on the gpu

This commit is contained in:
Dhairya Gandhi 2020-01-17 15:52:26 +05:30
parent b1e68813a8
commit 29ab410794
1 changed files with 20 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
test/cuda/layers.jl
View File

@ -8,20 +8,25 @@
@test_broken gradient(x -> sum(gpu(x)), rand(3,3)) isa Tuple @test_broken gradient(x -> sum(gpu(x)), rand(3,3)) isa Tuple
@test_throws ErrorException gradient(x -> sum(cpu(x)), gpu(rand(3,3))) isa Tuple @test_throws ErrorException gradient(x -> sum(cpu(x)), gpu(rand(3,3))) isa Tuple
function gradtest(layers, args...; name = "Conv", xs = rand(Float32, 28, 28, 1, 1)) function gradtest(name::String, layers::Vector, xs = nothing, args...)
isnothing(xs) && error("Missing input to test the layers against.")
@testset "$name GPU grad tests" begin @testset "$name GPU grad tests" begin
for layer in layers for layer in layers
@testset "$layer GPU grad test" begin @testset "$layer GPU grad test" begin
l = gpu(layer(args...)) l = gpu(layer(args...))
xs = gpu(xs) xs = gpu(xs)
if l isa DepthwiseConv if l isa DepthwiseConv
@test_broken gradient(Flux.params(l)) do ps = Flux.params(l)
sum(l(xs)) @test_broken gradient(() -> sum(l(xs)), ps) isa Flux.Zygote.Grads
end isa Flux.Zygote.Grads
else else
@test gradient(Flux.params(l)) do ps = Flux.params(l)
sum(l(xs)) @test gradient(() -> sum(l(xs)), ps) isa Flux.Zygote.Grads
end isa Flux.Zygote.Grads gs = gradient(() -> sum(l(xs)), ps)
# Handle pooling layers
if !isempty(ps)
@test gs[first(ps)] isa Flux.CuArrays.CuArray
end
end end
end end
end end
@ -30,27 +35,28 @@ end
# Repeats from Conv, CrossCor # Repeats from Conv, CrossCor
r = rand(Float32, 28, 28, 1, 1)
conv_layers = [Conv, ConvTranspose, CrossCor, DepthwiseConv] conv_layers = [Conv, ConvTranspose, CrossCor, DepthwiseConv]
gradtest(conv_layers, (2,2), 1=>3, name = "Conv") gradtest("Conv", conv_layers, r, (2,2), 1=>3)
pooling_layers = [MaxPool, MeanPool] pooling_layers = [MaxPool, MeanPool]
gradtest(pooling_layers, (2,2), name = "Pooling") gradtest("Pooling", pooling_layers, r, (2,2))
dropout_layers = [Dropout, AlphaDropout] dropout_layers = [Dropout, AlphaDropout]
gradtest(dropout_layers, 0.5f0, name = "Dropout") gradtest("Dropout", dropout_layers, r, 0.5f0)
norm_layers = [LayerNorm, BatchNorm] norm_layers = [LayerNorm, BatchNorm]
gradtest(norm_layers, 1, name = "Normalising", xs = rand(Float32, 28,28,3,1)) gradtest("Normalising" norm_layers, rand(Float32, 28,28,3,1), 1)
instancenorm = [InstanceNorm] instancenorm = [InstanceNorm]
gradtest(instancenorm, 1, name = "InstanceNorm") gradtest(instancenorm, r, "InstanceNorm", 1)
groupnorm = [GroupNorm] groupnorm = [GroupNorm]
gradtest(groupnorm, 3, 1, name = "GroupNorm", xs = rand(Float32, 28,28,3,1)) gradtest("GroupNorm", groupnorm, rand(Float32, 28,28,3,1), 3, 1)
const stateless_layers = [Flux.mse, const stateless_layers = [Flux.mse,
Flux.crossentropy, Flux.crossentropy,
Flux.logitcrossentropy,] Flux.logitcrossentropy,
Flux.normalise] Flux.normalise]
const stateless_layers_broadcasted = [Flux.binarycrossentropy, const stateless_layers_broadcasted = [Flux.binarycrossentropy,