I am wondering whether straight through gradient estimator for categorical variables is implemented. I want to see if I can implement the vq-vae architecture with Pyro. Thanks
We have an implementation for relaxed categorical straight through in pyro.distributions.relaxed_straight_through
on our current dev
branch. You should be able to use that for your purpose (and it will be great if you could contribute it as a Pyro example!). Let us know if you face any issues.
That’s great. I have modified the vae example with convolutional layers. Next up is the vq implementation. I might also leave this for students as part of a capstone project.