Excellent work (mae.ipynb)!

[sayakpaul]

@ariG23498 this is fantastic stuff. Super clean, readable, and coherent with the original implementation. A couple of suggestions that would likely make things even better:

• Since you have already implemented masking visualization utilities how about making them part of the PatchEncoder itself? That way you could let it accept a test image, apply random masking, and plot it just like the way you are doing in the earlier cells. This way I believe the notebook will be cleaner.
• AdamW (tfa.optimizers.adamw) is a better choice when it comes to training Transformer-based models.
• Are we taking the loss on the correct component? I remember you mentioning it being dealt with differently.

After these points are addressed I will take a crack at porting the training loop to TPUs along with other performance monitoring callbacks.

[ariG23498]

Since you have already implemented masking visualization utilities how about making them part of the PatchEncoder itself? That way you could let it accept a test image, apply random masking, and plot it just like the way you are doing in the earlier cells. This way I believe the notebook will be cleaner.

This makes sense. I will try incorporating this.

AdamW (tfa.optimizers.adamw) is a better choice when it comes to training Transformer-based models.

Noted!

Are we taking the loss on the correct component? I remember you mentioning it being dealt with differently.

Yes you are right! You will find the correct loss implementation in the mae_loss.ipynb notebook.

TODO:

• Incorporating AdamW
• Putting the visual utility inside PatchEncoder
• Get everything inside a single notebook
• Reuse shared logic inside the MAE class as suggested by @sayakpaul

[ariG23498]

After these points are addressed I will take a crack at porting the training loop to TPUs along with other performance monitoring callbacks.

Can't wait for this baby to train! 😃

[sayakpaul]

this happens as soon as you are okay with the mae_loss.ipynb implementation.

Taking a look now.

[sayakpaul]

@ariG23498 mae_loss.ipynb looks good. Since train_step() and test_step() share a large block of code, let's make another class method implementing that shared logic so that it can be reused. The method could be parameterized on a training argument to distinguish between training and inference.

[ariG23498]

I have updated the TODO accordingly.

[ariG23498]

@sayakpaul I have pushed a single notebook MaskedAutoEncoders.ipynb in this commit 7e8788a
and have deleted all the other notebook for clarity. All the todos are done. Please provide your feedback on the same.

[ariG23498]

Closing this.
#2 takes care of all the TODOs.