slimgroup/InvertibleNetworks.jl

Wavelet squeeze is not GPU friendly

yl4070 opened this issue · 3 comments

yl4070 commented

Hi, on my testing, wavelet_squeeze cause scalar indexing warning. And on benchmark, the it appears it would took almost 2s to complete such calculation. Is there any way to work around this?

X = rand(Float32, 16, 16, 16, 10) |> gpu
@btime wavelet_squeeze(X)
# 1.988s
mloubout commented

wavelet_squeeze uses an external Wavelet transform librairy that doesn't support GPU and shouldn't be used for GPU cases. If you want to use networks and layer on GPU involving squeeze you should use our Harr implementation that is a custom implementation of the Wavelet transform with a Harr basis that supports GPU. Check Haar_squeeze and invHaar_unsqueeze

yl4070 commented

Thanks. I wonder if this package is still in active development? Can I expect new feature/layer/network get implemented?

mloubout commented

This package is still very active yes. We haven't come across much new networks/layers to add but we are still working with this package daily for our research and updates and improvement will be added as much as we can.