fix ctc_gpu

FluxML · Jul 10, 2021 · 0be6765 · 0be6765
1 parent f28343f
commit 0be6765
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/src/losses/ctc-gpu.jl b/src/losses/ctc-gpu.jl
@@ -204,14 +204,15 @@ end
 function ctc_alpha(ŷ::CuArray, y)
   ŷ = logsoftmax(ŷ)
   blank = size(ŷ, 1)
-  z′ = fill(blank, 2 * length(y) + 1)
-  z′[eachindex(y) .* 2] = y
+  ycu = cu(y) 
+  z′ = CUDA.fill(blank, 2 * length(y) + 1)
+  z′[eachindex(y) .* 2] .= ycu
   T = size(ŷ, 2)
   U′ = 2*length(y) + 1
-  alphas = CUDA.fill(log(zero(ŷ[1])), U′,T)
-  nRepeats = count_repeats(y)
+  alphas = CUDA.fill(log(zero(eltype(ŷ))), U′,T)
+  nRepeats = count_repeats(cpu(y))
   nThreads = min(U′, MAX_THREADS)
-  @cuda blocks=1 threads=nThreads compute_alpha_kernel(ŷ, length(y), T, nRepeats, CuArray(y), CuArray(z′), alphas, blank)
+  @cuda blocks=1 threads=nThreads compute_alpha_kernel(ŷ, length(y), T, nRepeats, ycu, z′, alphas, blank)
   return (loss=-1 * logsumexp(alphas[end-1:end]), alpha=alphas, z′=z′, yhat=ŷ, nRepeats=nRepeats)
 end
 
@@ -221,7 +222,7 @@ function ∇ctc_loss(ŷ::CuArray, y, out)
   loss, alphas, z′, ŷ, nRepeats = out
   U′, T = size(alphas)
   blank = size(ŷ, 1)
-  typed_zero = zero(first(ŷ))
+  typed_zero = zero(eltype(ŷ))
   betas = CUDA.fill(log(typed_zero), U′, T)
   output = CUDA.fill(log(typed_zero), U′, T)
   nThreads = min(U′, MAX_THREADS)

diff --git a/test/outputsize.jl b/test/outputsize.jl
@@ -10,6 +10,7 @@
   @test outputsize(m, (10,); padbatch=true) == (2, 1)
   @test outputsize(m, (10, 30)) == (2, 30)
 
+  @info "Don't mind the following error, it's for testing purpose."
   m = Chain(Dense(10, 8, σ), Dense(8, 4), Dense(5, 2))
   @test_throws DimensionMismatch outputsize(m, (10,))
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,5 +1,6 @@
 using Flux
 using Flux.Data
+using Flux: OneHotArray, OneHotMatrix, OneHotVector
 using Test
 using Random, Statistics, LinearAlgebra
 using IterTools: ncycle
@@ -50,7 +51,7 @@ end
   end
 end
 
-@static if VERSION == v"1.5"
+@static if VERSION == v"1.6"
   using Documenter
   @testset "Docs" begin
     DocMeta.setdocmeta!(Flux, :DocTestSetup, :(using Flux); recursive=true)