Hi,
I’m just getting into Flux (and Julia more generally), but I’m finding it impossible to get rid of a very particular error in my convolutional VAE model. The code runs fine until the back(1f0) line where it throws ERROR: LoadError: UndefRefError: access to undefined reference.
I suspect this might the same core problem with Float32 being converted to Float64 (as documented here). I just cannot seem to resolve the issue even though my arrays seem to all be Float32 when I print them out.
Here is my minimal example (Julia: Version 1.4.1 (2020-04-14), Flux: Version 0.10.4):
using Flux
using Flux: logitbinarycrossentropy
using Flux.Data: DataLoader
using ImageFiltering
using MLDatasets: FashionMNIST
using Random
using Zygote
function get_train_loader(batch_size, shuffle::Bool)
# FashionMNIST is made up of 60k 28 by 28 greyscale images
train_x, train_y = FashionMNIST.traindata(Float32)
train_x = reshape(train_x, (28, 28, 1, :))
train_x = parent(padarray(train_x, Fill(0, (2,2,0,0))))
return DataLoader(train_x, train_y, batchsize=batch_size, shuffle=shuffle)
end
function vae_loss(encoder_μ, encoder_logσ, decoder, x)
batch_size = size(x)[end]
# Forward propagate through encoder
μ = encoder_μ(x)
logσ = encoder_logσ(x)
# Apply reparameterisation trick to sample latent
z = μ + randn(Float32, size(logσ)) .* exp.(logσ)
# Reconstruct from latent sample
x̂ = decoder(z)
# Negative reconstruction loss Ε_q[logp_x_z]
logp_x_z = -sum(logitbinarycrossentropy.(x̂, x)) / batch_size
# KL(qᵩ(z|x)||p(z)) where p(z)=N(0,1) and qᵩ(z|x) models the encoder i.e. reverse KL
# The @. macro makes sure that all operates are elementwise
kl_q_p = 0.5f0 * sum(@. (exp(2f0*logσ) + μ^2f0 - 2f0*logσ - 1f0)) / batch_size
# We want to maximise the evidence lower bound (ELBO)
β = 1f0
elbo = logp_x_z - β .* kl_q_p
return -elbo
end
function train()
# Define the encoder network
encoder_features = Chain(
Conv((4, 4), 1 => 32, relu; stride = 2, pad = 1),
Conv((4, 4), 32 => 32, relu; stride = 2, pad = 1),
Conv((4, 4), 32 => 32, relu; stride = 2, pad = 1),
flatten,
Dense(32 * 4 * 4, 256, relu),
)
encoder_μ = Chain(encoder_features, Dense(256, 10))
encoder_logσ = Chain(encoder_features, Dense(256, 10))
# Define the decoder network
decoder = Chain(
Dense(10, 256, relu),
Dense(256, 32 * 4 * 4, relu),
x -> reshape(x, (4, 4, 32, 16)),
ConvTranspose((4, 4), 32 => 32, relu; stride = 2, pad = 1),
ConvTranspose((4, 4), 32 => 32, relu; stride = 2, pad = 1),
ConvTranspose((4, 4), 32 => 1; stride = 2, pad = 1)
)
trainable_params = Flux.params(encoder_μ, encoder_logσ, decoder)
optimiser = ADAM(0.0001, (0.9, 0.999))
batch_size = 16
shuffle_data = true
dataloader = get_train_loader(batch_size, shuffle_data)
for epoch_num = 1:3
for (x_batch, y_batch) in dataloader
# pullback function returns the result (loss) and a pullback operator (back)
loss, back = Zygote.pullback(trainable_params) do
vae_loss(encoder_μ, encoder_logσ, decoder, x_batch)
end
# Feed the pullback 1 to obtain the gradients and update the model parameters
# NOTE: This is where it seems to fail
gradients = back(1f0)
Flux.Optimise.update!(optimiser, trainable_params, gradients)
println(loss)
end
end
println("Training complete!")
end
if abspath(PROGRAM_FILE) == @__FILE__
train()
end
The stack trace (and preceding warning):
┌ Warning: `haskey(::TargetIterator, name::String)` is deprecated, use `Target(; name = name) !== nothing` instead.
│ caller = llvm_compat(::VersionNumber) at compatibility.jl:176
└ @ CUDAnative ~/.julia/packages/CUDAnative/ierw8/src/compatibility.jl:176
ERROR: LoadError: UndefRefError: access to undefined reference
Stacktrace:
[1] getindex at ./array.jl:789 [inlined]
[2] conv_direct!(::Array{AbstractFloat,5}, ::Array{AbstractFloat,5}, ::Array{Float32,5}, ::DenseConvDims{3,(4, 4, 1),32,32,(2, 2, 1),(1, 1, 1, 1, 0, 0),(1, 1, 1),false}; alpha::Float64, beta::Bool) at /home/aleco/.julia/packages/NNlib/FAI3o/src/impl/conv_direct.jl:98
[3] conv_direct! at /home/aleco/.julia/packages/NNlib/FAI3o/src/impl/conv_direct.jl:51 [inlined]
[4] conv!(::Array{AbstractFloat,5}, ::Array{AbstractFloat,5}, ::Array{Float32,5}, ::DenseConvDims{3,(4, 4, 1),32,32,(2, 2, 1),(1, 1, 1, 1, 0, 0),(1, 1, 1),false}; kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /home/aleco/.julia/packages/NNlib/FAI3o/src/conv.jl:99
[5] conv!(::Array{AbstractFloat,5}, ::Array{AbstractFloat,5}, ::Array{Float32,5}, ::DenseConvDims{3,(4, 4, 1),32,32,(2, 2, 1),(1, 1, 1, 1, 0, 0),(1, 1, 1),false}) at /home/aleco/.julia/packages/NNlib/FAI3o/src/conv.jl:97
[6] conv!(::Array{AbstractFloat,4}, ::Array{AbstractFloat,4}, ::Array{Float32,4}, ::DenseConvDims{2,(4, 4),32,32,(2, 2),(1, 1, 1, 1),(1, 1),false}; kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /home/aleco/.julia/packages/NNlib/FAI3o/src/conv.jl:70
[7] conv! at /home/aleco/.julia/packages/NNlib/FAI3o/src/conv.jl:70 [inlined]
[8] conv(::Array{AbstractFloat,4}, ::Array{Float32,4}, ::DenseConvDims{2,(4, 4),32,32,(2, 2),(1, 1, 1, 1),(1, 1),false}; kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /home/aleco/.julia/packages/NNlib/FAI3o/src/conv.jl:116
[9] conv at /home/aleco/.julia/packages/NNlib/FAI3o/src/conv.jl:114 [inlined]
[10] #1224 at /home/aleco/.julia/packages/Zygote/1GXzF/src/lib/nnlib.jl:43 [inlined]
[11] (::Zygote.var"#2720#back#1226"{Zygote.var"#1224#1225"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Array{Float32,4},Array{Float32,4},DenseConvDims{2,(4, 4),32,32,(2, 2),(1, 1, 1, 1),(1, 1),false}}})(::Array{AbstractFloat,4}) at /home/aleco/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49
[12] ConvTranspose at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/conv.jl:148 [inlined]
[13] (::typeof(∂(λ)))(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
[14] applychain at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/basic.jl:36 [inlined]
[15] (::typeof(∂(applychain)))(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
[16] applychain at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/basic.jl:36 [inlined]
[17] (::typeof(∂(applychain)))(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
[18] applychain at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/basic.jl:36 [inlined]
[19] (::typeof(∂(applychain)))(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
[20] applychain at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/basic.jl:36 [inlined]
[21] (::typeof(∂(applychain)))(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
[22] applychain at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/basic.jl:36 [inlined]
[23] (::typeof(∂(applychain)))(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
[24] Chain at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/basic.jl:38 [inlined]
[25] (::typeof(∂(λ)))(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
[26] vae_loss at /home/aleco/Documents/flux-vae/minimal_example/online.jl:26 [inlined]
[27] (::typeof(∂(vae_loss)))(::Float32) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
[28] #4 at /home/aleco/Documents/flux-vae/minimal_example/online.jl:71 [inlined]
[29] (::typeof(∂(λ)))(::Float32) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
[30] (::Zygote.var"#50#51"{Params,Zygote.Context,typeof(∂(λ))})(::Float32) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface.jl:177
[31] train() at /home/aleco/Documents/flux-vae/minimal_example/online.jl:75
[32] top-level scope at /home/aleco/Documents/flux-vae/minimal_example/online.jl:85
[33] include(::Module, ::String) at ./Base.jl:377
[34] exec_options(::Base.JLOptions) at ./client.jl:288
[35] _start() at ./client.jl:484
in expression starting at /home/aleco/Documents/flux-vae/minimal_example/online.jl:84
Any help would be greatly appreciated as I am really having a hard time trying to debug this.
Kind regards,
Aleco
