I'm using Julia for UNet training with GPU. (Runs without issues on CPU). The error is at NNlib's function conv_im2col()
function conv_im2col!(
y::AbstractArray{T,5}, x::AbstractArray{T,5},
w::AbstractArray{T,5}, cdims::DenseConvDims;
col::AbstractArray{T,3}=similar(x, im2col_dims(cdims)),
alpha::T=T(1), beta::T=T(0)) where {T}
check_dims(size(x), size(w), size(y), cdims)
# COL * W -> Y
# [M x K] * [K x N] -> [M x N]
#
# M: output spatial resolution
# N: output channels
# K: size of input "patch" (kernel size and input channels combined)
#
# In english, we're grabbing each input patch and laying them out along
# the M dimension in `col`, so that the GEMM call below multiplies each
# kernel (which is kernel_h * kernel_w * channels_in elments long) is
# dotproducted with that input patch, effectively computing a convolution
# in a somewhat memory-wasteful but easily-computed way (since we already
# have an extremely highly-optimized GEMM call available in BLAS).
M = prod(output_size(cdims))
N = channels_out(cdims)
K = prod(kernel_size(cdims))*channels_in(cdims)
@threads for batch_idx in 1:size(x,5)
# col_slice is a thread-local workspace
col_slice = view(col, :, :, threadid())
im2col!(col_slice, view(x, :, :, :, :, batch_idx), cdims)
GC.@preserve col_slice w y begin
col_ptr = pointer(col_slice)
w_ptr = pointer(w)
y_ptr = pointer(y, (batch_idx - 1)*M*N + 1)
gemm!(Val(false), Val(false), M, N, K, alpha, col_ptr, w_ptr, beta, y_ptr)
end
and the detailed error is:
[ Info: training... [ Info: Using Device gpu Epoch Minibatch Loss
Test Dice Coef ---------------------------------------------- 1ERROR: LoadError: TaskFailedExceptionnested task error: TaskFailedException Stacktrace: [1] wait @ ./task.jl:322 [inlined] [2] threading_run(func::Function) @ Base.Threads ./threadingconstructs.jl:34 [3] macro expansion @ ./threadingconstructs.jl:93 [inlined] [4] conv_im2col!(y::SubArray{Float32, 5, Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false}, x::SubArray{Float32, 5, Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false}, w::CuArray{Float32, 5, CUDA.Mem.DeviceBuffer}, cdims::DenseConvDims{3, 3, 3, 6, 3}; col::Array{Float32, 3}, alpha::Float32, beta::Float32) @ NNlib ~/.julia/packages/NNlib/0QnJJ/src/impl/conv_im2col.jl:47
[5] conv_im2col! @ ~/.julia/packages/NNlib/0QnJJ/src/impl/conv_im2col.jl:28 [inlined]
[6] (::NNlib.var"#262#266"{Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, DenseConvDims{3, 3, 3, 6, 3},SubArray{Float32, 5, Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false}, CuArray{Float32, 5, CUDA.Mem.DeviceBuffer}, SubArray{Float32, 5, Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false}})() @ NNlib ./threadingconstructs.jl:169
nested task error: MethodError: no method matching gemm!(::Val{false}, ::Val{false}, ::Int64, ::Int64, ::Int64,::Float32, ::Ptr{Float32}, ::CuPtr{Float32}, ::Float32, ::Ptr{Float32}) Closest candidates are: gemm!(::Val, ::Val, ::Int64, ::Int64, ::Int64, ::Float32, ::Ptr{Float32}, ::Ptr{Float32}, ::Float32, ::Ptr{Float32}) at.julia/packages/NNlib/0QnJJ/src/gemm.jl:29 gemm!(::Val, ::Val, ::Int64, ::Int64, ::Int64, ::Float64, ::Ptr{Float64}, ::Ptr{Float64}, ::Float64, ::Ptr{Float64}) at .julia/packages/NNlib/0QnJJ/src/gemm.jl:29 gemm!(::Val, ::Val, ::Int64, ::Int64, ::Int64, ::ComplexF64, ::Ptr{ComplexF64}, ::Ptr{ComplexF64}, ::ComplexF64, ::Ptr{ComplexF64}) at .julia/packages/NNlib/0QnJJ/src/gemm.jl:29 ... Stacktrace: [1] macro expansion @ ~/.julia/packages/NNlib/0QnJJ/src/impl/conv_im2col.jl:56 [inlined] [2] (::NNlib.var"#909#threadsfor_fun#504"{Array{Float32, 3}, Float32, Float32, SubArray{Float32, 5, Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false}, SubArray{Float32, 5, Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false}, CuArray{Float32, 5, CUDA.Mem.DeviceBuffer}, DenseConvDims{3, 3, 3, 6, 3}, Int64, Int64, Int64, UnitRange{Int64}})(onethread::Bool) @ NNlib ./threadingconstructs.jl:81 [3] (::NNlib.var"#909#threadsfor_fun#504"{Array{Float32, 3}, Float32, Float32, SubArray{Float32, 5, Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false}, SubArray{Float32, 5, Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false}, CuArray{Float32, 5, CUDA.Mem.DeviceBuffer}, DenseConvDims{3, 3, 3, 6, 3}, Int64, Int64, Int64, UnitRange{Int64}})() @ NNlib ./threadingconstructs.jl:48 Stacktrace: [1] sync_end(c::Channel{Any}) @ Base ./task.jl:369 [2] macro expansion @ ./task.jl:388 [inlined] [3] conv!(out::Array{Float32, 5}, in1::Array{Float32, 5}, in2::CuArray{Float32, 5, CUDA.Mem.DeviceBuffer}, cdims::DenseConvDims{3, 3, 3, 6, 3}; kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) @ NNlib ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:195 [4] conv! @ ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:186 [inlined] [5] conv!(y::Array{Float32, 4}, x::Array{Float32, 4}, w::CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, cdims::DenseConvDims{2, 2, 2, 4, 2}; kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) @ NNlib ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:145 [6] conv! @ ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:145 [inlined] [7] conv(x::Array{Float32, 4}, w::CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, cdims::DenseConvDims{2, 2, 2, 4, 2}; kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) @ NNlib ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:88 [8] conv @ ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:86 [inlined] [9] #rrule#312 @ ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:313 [inlined] [10] rrule @ ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:304 [inlined] [11] rrule @ ~/.julia/packages/ChainRulesCore/ctmSK/src/rules.jl:134 [inlined] [12] chain_rrule @ ~/.julia/packages/Zygote/xGkZ5/src/compiler/chainrules.jl:218 [inlined] [13] macro expansion @ ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [inlined] [14] _pullback @ ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:9 [inlined] [15] _pullback @ ~/.julia/packages/Flux/KkC79/src/layers/conv.jl:200 [inlined] [16] macro expansion @ ~/.julia/packages/Flux/KkC79/src/layers/basic.jl:53 [inlined] [17] _pullback @ ~/.julia/packages/Flux/KkC79/src/layers/basic.jl:53 [inlined] [18] _pullback(::Zygote.Context{true}, ::typeof(Flux._applychain), ::WARNING: both Losses and NNlib export "ctc_loss"; uses of it in module Flux must be qualified Tuple{Conv{2, 2, typeof(identity), CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(selu), CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Conv{2, 2, typeof(identity), CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(selu), CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}, ::Array{Float32, 4}) @ Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [19] _pullback @ ~/.julia/packages/Flux/KkC79/src/layers/basic.jl:51 [inlined] [20] _pullback(ctx::Zygote.Context{true}, f::Chain{Tuple{Conv{2, 2, typeof(identity), CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(selu), CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Conv{2, 2, typeof(identity), CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(selu), CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, args::Array{Float32, 4}) @ Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [21] _pullback @ ~/.julia/dev/UNet/src/model.jl:93 [inlined] [22] _pullback(ctx::Zygote.Context{true}, f::Unet, args::Array{Float32, 4}) @ Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [23] _pullback @ ~/.julia/dev/UNet/src/train.jl:37 [inlined] [24] _pullback(::Zygote.Context{true}, ::UNet.var"#loss#75"{typeof(Flux.Losses.logitcrossentropy), Unet}, ::Array{Float32, 4}, ::Array{Int32, 4}) @ Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [25] _apply @ ./boot.jl:804 [inlined] [26] adjoint @ ~/.julia/packages/Zygote/xGkZ5/src/lib/lib.jl:203 [inlined] [27] _pullback @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined] [28] _pullback @ ~/.julia/packages/Flux/KkC79/src/optimise/train.jl:120 [inlined] [29] _pullback(::Zygote.Context{true}, ::Flux.Optimise.var"#37#40"{UNet.var"#loss#75"{typeof(Flux.Losses.logitcrossentropy), Unet}, Tuple{Array{Float32, 4}, Array{Int32, 4}}}) @ Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [30] pullback(f::Function, ps::Zygote.Params{Zygote.Buffer{Any, Vector{Any}}}) @ Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface.jl:373 [31] gradient(f::Function, args::Zygote.Params{Zygote.Buffer{Any, Vector{Any}}}) @ Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface.jl:96 [32] macro expansion @ ~/.julia/packages/Flux/KkC79/src/optimise/train.jl:119 [inlined] [33] macro expansion @ ~/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328 [inlined] [34] train!(loss::Function, ps::Zygote.Params{Zygote.Buffer{Any, Vector{Any}}}, data::MLUtils.DataLoader{Tuple{Array{Float32, 4}, Array{Int32, 4}}, Random._GLOBAL_RNG}, opt::RMSProp; cb::UNet.var"#77#78"{Unet, MLUtils.DataLoader{Tuple{Array{Float32, 4}, Array{Int32, 4}}, Random._GLOBAL_RNG}}) @ Flux.Optimise ~/.julia/packages/Flux/KkC79/src/optimise/train.jl:117 [35] train(train_dataset::ImageDataset, test_dataset::ImageDataset, model::Unet, opt::RMSProp, loss::UNet.var"#loss#75"{typeof(Flux.Losses.logitcrossentropy), Unet}; batch_size::Int64, num_epochs::Int64, learning_rate_drop_rate::Float64, learning_rate_step::Int64, device::Function, save_intermediate_models::Bool, save_model::Bool, out_dir::String, model_file_name::String) @ UNet ~/.julia/dev/UNet/src/train.jl:123 [36] train(train_dataset::ImageDataset, test_dataset::ImageDataset, num_channels::Int64, num_labels::Int64; loss_function::typeof(Flux.Losses.logitcrossentropy), learning_rate::Float64, momentum::Float64, device::Function, kwargs::Base.Iterators.Pairs{Symbol, Any, NTuple{6, Symbol}, NamedTuple{(:num_epochs, :batch_size, :save_intermediate_models, :save_model, :out_dir, :model_file_name), Tuple{Int64, Int64, Bool, Bool, String, String}}}) @ UNet ~/.julia/dev/UNet/src/train.jl:40 [37] top-level scope @ /train_images.jl:70 [38] include(fname::String) @ Base.MainInclude ./client.jl:444 [39] top-level scope @ REPL[2]:1 in expression starting at train_images.jl:70