17 gpu acceleration fails for mljflux

Project.toml

@@ -17,6 +17,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [compat]
 CategoricalArrays = "0.10"

docs/Project.toml

@@ -16,3 +16,4 @@ ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"

docs/setup_docs.jl

@@ -7,7 +7,7 @@ setup_docs = quote
     using Distributions
     using Flux
     using Flux.Data: DataLoader
-    using Flux: onehotbatch, onecold, @epochs, @functor, logsumexp
+    using Flux: onehotbatch, onecold, @functor, logsumexp
     using Flux.Losses: logitcrossentropy
     using Images
     using JointEnergyModels

src/Samplers/Samplers.jl

@@ -88,7 +88,7 @@ function ConditionalSampler(
     max_len::Int=10000, prob_buffer::AbstractFloat=0.95
 )
     @assert batch_size <= max_len "batch_size must be <= max_len"
-    buffer = Float32.(rand(𝒟x, input_size..., batch_size))
+    buffer = Float32.(rand(𝒟x, input_size..., batch_size)) 
     return ConditionalSampler(𝒟x, 𝒟y, input_size, batch_size, buffer, max_len, prob_buffer)
 end
 

src/mlj_flux.jl

@@ -121,6 +121,11 @@ function MLJFlux.fit!(model::JointEnergyClassifier, penalty, chain, optimiser, e
     train_set = zip(X, y)
     opt_state = Flux.setup(optimiser, model.jem)
 
+    # move sampler to GPU if to_gpu:
+    if model.acceleration isa CUDALibs
+        model.sampler = gpu(model.sampler)
+    end
+
     history = train_model(
         model.jem, train_set, opt_state;
         class_loss_fun=loss,

src/model.jl

@@ -51,7 +51,12 @@ function get_samples(jem::JointEnergyModel, x)
         ndims(jem.sampler.buffer), 
         1:size_sample
     )
+
     @assert size(xsample) == size(x)
+
+    @info "Using GPU for sampling."
+    x, xsample = gpu(x), gpu(xsample)
+
     return x, xsample
 end
 

src/training.jl

@@ -53,6 +53,8 @@ function train_model(
     class_loss_fun::Function=logitcrossentropy,
     progress_meter::Union{Nothing,ProgressMeter.Progress}=nothing,
 )
+
+    # Set up:
     training_log = []
     not_finite_counter = 0
     if isnothing(progress_meter)
@@ -63,6 +65,7 @@ function train_model(
         verbosity == 0 || next!(progress_meter)
     end
 
+    # Training loop:
     for epoch in 1:num_epochs
         training_losses = Float32[]