Differentiable Regularizers for CUDA, CPUs?

Specific Domains
1

Hi,

is anyone aware of a package which implements Total Variation (TV) or other kinds of regularizer efficiently and differentiable?

A naive, differentiable way of TV would be:

julia> arr = rand(10, 10)
10×10 Matrix{Float64}:
 0.91845   0.141477  0.150494  0.832751  0.436746  0.402896   0.760819    0.123788   0.701573   0.828201
 0.807974  0.672054  0.632665  0.247365  0.565101  0.768765   0.140932    0.673522   0.383663   0.496613
 0.796742  0.573989  0.718057  0.801975  0.348846  0.966668   0.395765    0.0118382  0.234899   0.4305
 0.270457  0.235942  0.302115  0.159628  0.642272  0.0840647  0.298663    0.0900039  0.324911   0.422608
 0.193854  0.853308  0.271397  0.675382  0.841182  0.584587   0.30056     0.744804   0.774442   0.311683
 0.987725  0.769463  0.286531  0.297112  0.308935  0.0269483  0.825088    0.836484   0.802654   0.474037
 0.674593  0.709923  0.356206  0.360819  0.510561  0.109935   0.978717    0.471597   0.0912487  0.713041
 0.615551  0.128082  0.17853   0.855967  0.100066  0.237995   0.282599    0.503759   0.596497   0.927304
 0.95856   0.61049   0.859611  0.653335  0.581014  0.509047   0.00609067  0.911455   0.865939   0.879677
 0.511518  0.052382  0.704612  0.620247  0.294648  0.450548   0.178217    0.316866   0.968324   0.630102

julia> function TV(arr::AbstractArray{T}, ϵ=T(1e-8)) where T
           return sum(sqrt.(abs2.(arr .- circshift(arr, (1,0))) .+ abs2.(arr .- circshift(arr, (0,1))) .+ ϵ))
       end
TV (generic function with 2 methods)

julia> TV(arr)
47.522025415255186

Of course this creates many intermediate copies and is not very performant. Tullio.jl alllows to write such expressions incredibly fast on CPUs, on GPUs it’s not competitive, especially the gradients are slow.

I was wondering whether someone tried to tackle this with KernelAbstractions.jl or similar.

Best,

Felix

CC: @RainerHeintzmann

1 Like