Documentation Suggestion #233
-
|
Hi guys, I'm trying to implement a WGAN-GP in flax and got it to work but spent a few hours hunting for a bug that turned out to be me misunderstanding the library. So I figured I'd write you a note so you can include this gotcha in the documentation. The problem was that when I applied gradients to an optimizer, I'd assumed that my original copy of the model (in this case the Generator) could be called on some inputs and would contain the updated weights. Only at the end of the MNIST example when it says "At any particular optimzation step, This behavior seems very functional, which is fine, but certainly different than what I'd expect coming from PyTorch or TF. It would be great if that was made more clear in the documentation. Maybe even the optimizer / train step interface might be changed in the example to make it clear that Thanks, |
Beta Was this translation helpful? Give feedback.
Replies: 3 comments 1 reply
-
|
oh also I would be happy to contribute the code for the WGAN-GP as an example if that would be useful to you guys. I am pretty happy with the quality of the code. |
Beta Was this translation helpful? Give feedback.
-
|
Hi @virajmehta, nice to meet you! First of all, super awesome that you're using Flax to implement GANs -- we would love to link to your example. Per our new example policy, we avoid adding new official examples because of the long term maintenance expectations but we definitely want to highlight your model by linking to it prominently. Take a look at Per your comment on the confusion on the functional API, I think you're hitting a good point here. Our |
Beta Was this translation helpful? Give feedback.
-
|
Sounds great! I'll figure out how best to open source my code and will add a PR when I've figured that out. On the functional API comment, I think perhaps the train step of the MNIST example could make it clear either in comments or in code structure that the "live" model is only in optimizer.target. For example, my GAN and other code all returns optimizer and explicitly separately optimizer.target so that I can overwrite self.generator or self.critic or whatever in my training class, and I think that would pair nicely with an immutable Model. I don't think it's necessary to go that far, but I think an update there would have spared me the confusion. Thanks! |
Beta Was this translation helpful? Give feedback.
-
|
Yes, we are reworking our module API (inspired but not identical to the proposal in #208) and all examples and will probably make the top-level training loops hold a variables dict directly rather than a Model. I think that would solve at least part of the confusion you were hitting. |
Beta Was this translation helpful? Give feedback.
Hi @virajmehta, nice to meet you!
First of all, super awesome that you're using Flax to implement GANs -- we would love to link to your example. Per our new example policy, we avoid adding new official examples because of the long term maintenance expectations but we definitely want to highlight your model by linking to it prominently. Take a look at
flax/examples/README.mdand please do file a pull request adding a link to your code!Per your comment on the confusion on the functional API, I think you're hitting a good point here. Our
Modelwrapper is a convenience that also leads to users not being fully aware of the functional approach. One proposal is to just scrapModelaltogether, o…