arnauqb / GumbelSoftmax.jl

Julia implementation of the Gumbel-Softmax reparametrization trick compatible with Zygote and ForwardDiff
MIT License
6 stars 0 forks source link

Help with VAE and gumbel-softmax trick #4

Open flaviobarros opened 4 months ago

flaviobarros commented 4 months ago

Thanks for the code! I'm trying to learn Julia and Flux.jl but I'm having trouble finding an example of a VAE with gumbel-softmax trick. Chances are that you have already done this. Could you provide an example for me?

arnauqb commented 4 months ago

Hey, yes, I have an old one which I will try to clean up and rescue. I think it would be a good example for the repo! I should have a bit of time over the next 1-2 weeks to upload it.

arnauqb commented 3 months ago

As promised, I have added a discrete VAE example: https://github.com/arnauqb/GumbelSoftmax.jl/blob/main/examples/vae.jl

I've also taken the time to merge some changes I had locally (particularly the implementation of the Rao-Blackwell Gumbel-Softmax). Let me know if you find problems with the VAE implementation!