# Moving a custom loss function to GPU

**URL:** https://discourse.julialang.org/t/moving-a-custom-loss-function-to-gpu/109637
**Category:** Machine Learning
**Tags:** gpu, flux, zygote
**Created:** [February 2, 2024, 5:50pm UTC](https://discourse.julialang.org/t/moving-a-custom-loss-function-to-gpu/109637 "2024-02-02T17:50:46Z")
**Posts on this page:** 6
**Page:** 1

<div class="post-metadata">

### Author: ![josemanuel22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/josemanuel22/32/20668_2.png) [@josemanuel22](https://discourse.julialang.org/u/josemanuel22)
#### Post date: [February 2, 2024, 5:50pm UTC](https://discourse.julialang.org/t/moving-a-custom-loss-function-to-gpu/109637/1 "2024-02-02T17:50:46Z")

</div>

I have developed the following custom loss function, I am trying to run it on the GPU efficiently but I am unable to do so; I get the following error in the line `Matrix(ω' * yₖ_batch)`

```julia
ERROR: MethodError: no method matching unsafe_convert(::Type{Ptr{Float32}}, ::CUDA.CuPtr{Float32})
Closest candidates are:
  unsafe_convert(::Type{CUDA.CuRef{T}}, ::Any) where T at ~/.julia/packages/CUDA/35NC6/src/pointer.jl:208
  unsafe_convert(::Type{CUDA.PtrOrCuPtr{T}}, ::Any) where T at ~/.julia/packages/CUDA/35NC6/src/pointer.jl:118
  unsafe_convert(::Type{<:Union{CUDA.CuArrayPtr, CUDA.CuPtr, Ptr}}, ::CUDA.Mem.AbstractBuffer) at ~/.julia/packages/CUDA/35NC6/lib/cudadrv/memory.jl:33

```

The lost function is the following one,

```julia
function sliced_invariant_statistical_loss_optimized_2(nn_model, loader, hparams)
    @assert loader.batchsize == hparams.samples
    @assert length(loader) == hparams.epochs
    losses = Vector{Float32}()
    optim = Flux.setup(Flux.Adam(hparams.η), nn_model)

    @showprogress for data in loader
        Ω = gpu(ThreadsX.map(_ -> sample_random_direction(size(data)[1]), 1:(hparams.m)))
        loss, grads = Flux.withgradient(nn_model) do nn
            total = 0.0f0
            # Generate all random numbers in one go
            x_batch = gpu(rand(hparams.noise_model, hparams.samples * hparams.K))

            # Process batch through nn_model
            yₖ_batch = nn(Float32.(x_batch))
            for ω in Ω
                aₖ = zeros(Float32, hparams.K + 1) # Reset aₖ for each new ω

                s = Matrix(ω' * yₖ_batch)

                # Pre-compute column indices for slicing
                start_cols = hparams.K * (1:(hparams.samples - 1))
                end_cols = hparams.K * (2:(hparams.samples)) .- 1

                # Create slices of 's' for all 'aₖ_slice'
                aₖ_slices = [
                    s[:, start_col:(end_col - 1)] for
                    (start_col, end_col) in zip(start_cols, end_cols)
                ]

                # Compute the dot products for all iterations at once
                ω_data_dot_products = [dot(ω, data[:, i]) for i in 2:(hparams.samples)]

                # Apply 'generate_aₖ' for each pair and sum the results
                aₖ = sum([
                    generate_aₖ(aₖ_slice, ω_data_dot_product) for
                    (aₖ_slice, ω_data_dot_product) in zip(aₖ_slices, ω_data_dot_products)
                ])
                total += scalar_diff(aₖ ./ sum(aₖ))
            end
            total / hparams.m
        end
        Flux.update!(optim, nn_model, grads[1])
        push!(losses, loss)
    end
    return losses
end

```

And I call it this way,

```julia
device = gpu

model = device(Generator(latent_dim))
#model = Chain( ConvTranspose((7, 7), 100 => 256, stride=1, padding=0), BatchNorm(256, relu), ConvTranspose((4, 4), 256 => 128, stride=2, padding=1), BatchNorm(128, relu), ConvTranspose((4, 4), 128 => 1, stride=2, padding=1), tanh ))

# Mean vector (zero vector of length dim)
mean_vector = zeros(dims)

# Covariance matrix (identity matrix of size dim x dim)
cov_matrix = Diagonal(ones(dims))

# Create the multivariate normal distribution
noise_model = device(MvNormal(mean_vector, cov_matrix))

n_samples = 10000

hparams = device(
    HyperParamsSlicedISL(;
        K=10, samples=1000, epochs=60, η=1e-2, noise_model=noise_model, m=100
    ),
)

# Create a data loader for training
batch_size = 1000
#train_loader = DataLoader(train_x; batchsize=batch_size, shuffle=false, partial=false)
train_loader = gpu(DataLoader(train_x; batchsize=batch_size, shuffle=true, partial=false))

sliced_invariant_statistical_loss_optimized_2(model, train_loader, hparams)

```

can someone help me?

---

<div class="post-metadata">

### Author: ![maleadt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maleadt/32/10097_2.png) [@maleadt](https://discourse.julialang.org/u/maleadt)
#### Post date: [February 2, 2024, 8:38pm UTC](https://discourse.julialang.org/t/moving-a-custom-loss-function-to-gpu/109637/2 "2024-02-02T20:38:22Z")

</div>

> [@josemanuel22](#):
>
> `no method matching unsafe_convert(::Type{Ptr{Float32}}, ::CUDA.CuPtr{Float32})`

Without looking into the details, the above error suggests that you’re invoking CPU functionality assuming a CPU pointer with a GPU input (assuming the pointers were obtained by calling the `pointer` function). The backtrace should reveal more, but you didn’t include that in your post.

---

<div class="post-metadata">

### Author: ![josemanuel22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/josemanuel22/32/20668_2.png) [@josemanuel22](https://discourse.julialang.org/u/josemanuel22)
#### Post date: [February 2, 2024, 8:55pm UTC](https://discourse.julialang.org/t/moving-a-custom-loss-function-to-gpu/109637/3 "2024-02-02T20:55:04Z")

</div>

Thank you very much for your answer! The all backtrace is,

```julia
ERROR: MethodError: no method matching unsafe_convert(::Type{Ptr{Float32}}, ::CUDA.CuPtr{Float32})
Closest candidates are:
  unsafe_convert(::Type{CUDA.CuRef{T}}, ::Any) where T at ~/.julia/packages/CUDA/35NC6/src/pointer.jl:208
  unsafe_convert(::Type{CUDA.PtrOrCuPtr{T}}, ::Any) where T at ~/.julia/packages/CUDA/35NC6/src/pointer.jl:118
  unsafe_convert(::Type{<:Union{CUDA.CuArrayPtr, CUDA.CuPtr, Ptr}}, ::CUDA.Mem.AbstractBuffer) at ~/.julia/packages/CUDA/35NC6/lib/cudadrv/memory.jl:33
  ...
Stacktrace:
  [1] gemv!(trans::Char, alpha::Float32, A::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, X::Vector{Float32}, beta::Float32, Y::Vector{Float32})
    @ LinearAlgebra.BLAS /usr/share/julia/stdlib/v1.8/LinearAlgebra/src/blas.jl:666
  [2] gemv!(y::Vector{Float32}, tA::Char, A::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, x::Vector{Float32}, α::Bool, β::Bool)
    @ LinearAlgebra /usr/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:503
  [3] mul!
    @ /usr/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:65 [inlined]
  [4] mul!
    @ /usr/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:276 [inlined]
  [5] *
    @ /usr/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:52 [inlined]
  [6] *
    @ /usr/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:119 [inlined]
  [7] (::ChainRules.var"#1457#1460"{Adjoint{Float32, Vector{Float32}}, CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, ChainRulesCore.ProjectTo{Adjoint, NamedTuple{(:parent,), Tuple{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:element, :axes), Tuple{ChainRulesCore.ProjectTo{Float32, NamedTuple{(), Tuple{}}}, Tuple{Base.OneTo{Int64}}}}}}}}})()
    @ ChainRules ~/.julia/packages/ChainRules/pEOSw/src/rulesets/Base/arraymath.jl:36
  [8] unthunk
    @ ~/.julia/packages/ChainRulesCore/zoCjl/src/tangent_types/thunks.jl:204 [inlined]
  [9] wrap_chainrules_output
    @ ~/.julia/packages/Zygote/WOy6z/src/compiler/chainrules.jl:110 [inlined]
 [10] map
    @ ./tuple.jl:223 [inlined]
 [11] wrap_chainrules_output
    @ ~/.julia/packages/Zygote/WOy6z/src/compiler/chainrules.jl:111 [inlined]
 [12] (::Zygote.ZBack{ChainRules.var"#times_pullback#1459"{Adjoint{Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:element, :axes), Tuple{ChainRulesCore.ProjectTo{Float32, NamedTuple{(), Tuple{}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}}, ChainRulesCore.ProjectTo{Adjoint, NamedTuple{(:parent,), Tuple{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:element, :axes), Tuple{ChainRulesCore.ProjectTo{Float32, NamedTuple{(), Tuple{}}}, Tuple{Base.OneTo{Int64}}}}}}}}}})(dy::Adjoint{Float32, Vector{Float32}})
    @ Zygote ~/.julia/packages/Zygote/WOy6z/src/compiler/chainrules.jl:211
 [13] Pullback
    @ ~/Datos/github/ISL/src/CustomLossFunction.jl:667 [inlined]
 [14] (::Zygote.Pullback{Tuple{ISL.var"#376#381"{HyperParamsSlicedISL, CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, Vector{CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}, Chain{Tuple{Dense{typeof(identity), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, var"#5#7", ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, typeof(Flux.flatten), var"#6#8"}}}, Any})(Δ::Float32)
    @ Zygote ~/.julia/packages/Zygote/WOy6z/src/compiler/interface2.jl:0
 [15] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{ISL.var"#376#381"{HyperParamsSlicedISL, CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, Vector{CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}, Chain{Tuple{Dense{typeof(identity), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, var"#5#7", ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, typeof(Flux.flatten), var"#6#8"}}}, Any}})(Δ::Float32)
    @ Zygote ~/.julia/packages/Zygote/WOy6z/src/compiler/interface.jl:45
 [16] withgradient(f::Function, args::Chain{Tuple{Dense{typeof(identity), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, var"#5#7", ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, typeof(Flux.flatten), var"#6#8"}})
    @ Zygote ~/.julia/packages/Zygote/WOy6z/src/compiler/interface.jl:162
 [17] macro expansion
    @ ~/Datos/github/ISL/src/CustomLossFunction.jl:657 [inlined]
 [18] macro expansion
    @ ~/.julia/packages/ProgressMeter/vnCY0/src/ProgressMeter.jl:957 [inlined]
 [19] sliced_invariant_statistical_loss_optimized_2(nn_model::Chain{Tuple{Dense{typeof(identity), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, var"#5#7", ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(relu), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, ConvTranspose{2, 4, typeof(identity), CUDA.CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, typeof(Flux.flatten), var"#6#8"}}, loader::DataLoader{MLUtils.MappedData{:auto, typeof(gpu), Matrix{Float32}}, Random._GLOBAL_RNG, Val{nothing}}, hparams::HyperParamsSlicedISL)
    @ ISL ~/Datos/github/ISL/src/CustomLossFunction.jl:655
 [20] macro expansion
    @ ~/Datos/github/ISL/examples/Sliced_ISL/MNIST_sliced.jl:166 [inlined]
 [21] top-level scope
    @ ~/.julia/packages/ProgressMeter/vnCY0/src/ProgressMeter.jl:957
```

---

<div class="post-metadata">

### Author: ![maleadt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maleadt/32/10097_2.png) [@maleadt](https://discourse.julialang.org/u/maleadt)
#### Post date: [February 6, 2024, 8:04pm UTC](https://discourse.julialang.org/t/moving-a-custom-loss-function-to-gpu/109637/4 "2024-02-06T20:04:00Z")

</div>

That reveals a lot. The core problem is that you’re mixing a CPU and GPU operation:

```julia
gemv!(y::Vector{Float32}, tA::Char, A::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, x::Vector{Float32}, α::Bool, β::Bool)

```

Notice the CPU output, and a single GPU input. That results in CPU BLAS being used with a GPU buffer, which is unsupported (unless you use unified memory, which currently isn’t the default).

I’m not sure where the CPU arguments are introduced, but maybe this helps you figuring out the root of the issue.

---

<div class="post-metadata">

### Author: ![josemanuel22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/josemanuel22/32/20668_2.png) [@josemanuel22](https://discourse.julialang.org/u/josemanuel22)
#### Post date: [February 15, 2024, 9:17pm UTC](https://discourse.julialang.org/t/moving-a-custom-loss-function-to-gpu/109637/5 "2024-02-15T21:17:39Z")

</div>

Hello everyone, again. Thank you very much for the help. I think in the end I have located where I am mixing up GPU with CPU. And I believe it is here in this call, Is in this line,

```julia
aₖ = sum([ 
  generate_aₖ(aₖ_slice, ω_data_dot_product) for 
    (aₖ_slice, ω_data_dot_product) in zip(aₖ_slices, ω_data_dot_products)
  ])

```

To be more precise, the function `generate_aₖ` is,

```julia
function generate_aₖ(ŷ::CuArray{T}, y::T) where {T<:AbstractFloat}
    return CUDA.sum([γ(ŷ, y, k) for k in 0:length(ŷ)])
end

```

and the rest of the functions are,

```julia
function γ(yₖ::CuMatrix{T}, yₙ::T, m::Int64) where {T<:AbstractFloat}
    function eₘ_cuda(m, length)
        return CuArray([j == m ? T(1.0) : T(0.0) for j in 0:length])
    end

    return eₘ_cuda(m, size(yₖ, 2)) * ψₘ(ϕ(yₖ, yₙ), m) 
end

function ϕ(yₖ::CuMatrix{T}, yₙ::T) where {T<:AbstractFloat}
    return sum(_sigmoid(yₖ, yₙ))
end

function ψₘ(y::CuArray{T}, m::Int64) where {T<:AbstractFloat}
    stddev = T(0.1)
    return exp.(-0.5f0 * ((y .- m) / stddev) .^ 2)
end

function _sigmoid(ŷ::CuArray{T}, y::T) where {T<:AbstractFloat}
    return sigmoid_fast.((y .- ŷ) .* 10.0f0)
end

```

I think I can’t do a list comprehension in CUDA `CUDA.sum([γ(ŷ, y, k) for k in 0:length(ŷ)])`. Can someone tell me how to rewrite this or if this is where the problem is? Thank you very much.

---

<div class="post-metadata">

### Author: ![maleadt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maleadt/32/10097_2.png) [@maleadt](https://discourse.julialang.org/u/maleadt)
#### Post date: [February 15, 2024, 9:19pm UTC](https://discourse.julialang.org/t/moving-a-custom-loss-function-to-gpu/109637/6 "2024-02-15T21:19:49Z")

</div>

> [@josemanuel22](#):
>
> I think I can’t do a list comprehension in CUDA `CUDA.sum([γ(ŷ, y, k) for k in 0:length(ŷ)])`.

You can pass a `map` function to `sum` (which is basically `mapreduce(identity, +)`, in case you need more flexibility).