search

slimgroup / InvertibleNetworks.jl

A Julia framework for invertible neural networks
MIT License
149 stars 21 forks source link

Wavelet squeeze is not GPU friendly #52

Closed yl4070 closed 2 months ago

yl4070 commented 2 years ago

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 2 years ago

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 2 years ago

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

mloubout commented 2 years ago

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.