Add an Example: pix2pix

[us]

TODO:

• : Input pipeline
• : Generator and Discriminator Model
• : Calculate gradient
• : Jupyter notebook

I use TF pix2pix tutorial resource as a guide and I am stuck when calculating the gradient on train_step function. Can you help me?
Thank you!

[AlexeyG]

Hi @us

I am not familiar with the pix2pix model, so not sure if I can be of much help he, but I can try. Can you be more specific with what exactly you're stuck with?

Are you envisioning this as an example that would become a part of the canonical Flax examples? If so, please check our guidelines regarding adding new examples in examples/README.md.

[us]

Thank you! I read that and opened the issue #193.

I need to get grad the loss function but it's the correct way of that? I duplicated code.
Loss function of pix2pix needs to two model and when getting the get grad of loss_fn, is it correct way? I am so confused.
How can I get grads of two models? please check this for understanding loss func and grad of pix2pix.
I hope you understand this.

[avital]

This won't work as-is. You should have a single optimizer (because the optimizer keeps track of the model parameters). As you have it here, you're really not even optimizing the same model with these two calls to apply_gradient.

[avital]

Given no response in the past 26 days, I'll close this now. If you get a fully working version of pix2pix, we'd be happy to link to it form https://github.com/google/flax/blob/master/examples/README.md, per our new policy of preferring linking to external repos rather than maintaining them as part of Flax core.