feat: add BEiT support#1958
Conversation
- Refine the BEiT architecture - Add examples and notebooks - Expand unit tests - Improve documentation
|
/review |
|
Thanks a lot @aymen-000, what a cool contribution! This will probably take a bit of time to review, but we will try to do it as quickly as possible. |
| # NOTE: The ImageTokenizer must be pre-trained before BEIT pre-training. | ||
| # The tokenizer follows the DALL-E discrete VAE architecture and can be | ||
| # trained as follows: | ||
| # | ||
| # tokenizer = ImageTokenizer(vocab_size=8192) | ||
| # optimizer = torch.optim.Adam(tokenizer.parameters(), lr=1e-3) | ||
| # | ||
| # for images in dataloader: | ||
| # logits, recon = tokenizer(images) | ||
| # loss = F.mse_loss(recon, images) # Reconstruction loss | ||
| # loss.backward() | ||
| # optimizer.step() | ||
| # optimizer.zero_grad() | ||
| # | ||
| # torch.save(tokenizer.state_dict(), "tokenizer.pth") |
There was a problem hiding this comment.
I doubt that it's quite that easy (but unfortunately DALL-E does not have a code release and neither has Microsoft for the tokenize pretraining). Since I think that this is unlikely to yield any good results, I would rather remove it.
There was a problem hiding this comment.
Note: I found that BEiT seems to use a smooth L1 loss https://github.com/microsoft/unilm/blob/master/beit/modeling_discrete_vae.py#L124
But we don't have a training recipe, so even with that piece of information it's unlikely to get anything good.
| @@ -0,0 +1,187 @@ | |||
| """BEiT: BERT Pre-Training of Image Transformers. | |||
There was a problem hiding this comment.
We don't actually want these model modules anymore. Instead we want the full implementations to be inside the notebooks/examples. Sorry if this is not clear from the contributing instructions, will try to make this better in the future.
| BEiT | ||
| ==== | ||
|
|
||
| implementation of the BEiT (BERT Pre-Training of Image Transformers) |
There was a problem hiding this comment.
| implementation of the BEiT (BERT Pre-Training of Image Transformers) | |
| Implementation of the BEiT (BERT Pre-Training of Image Transformers) |
| # Copyright (c) 2021. Lightly AG and its affiliates. | ||
| # All Rights Reserved | ||
|
|
||
| from lightly.models.modules.beit_tokenizer import ImageTokenizer |
There was a problem hiding this comment.
I would call it BEiTImageTokenizer, so that it's clear to which pipeline it belongs. And I mean not only the import but the actual class.
| DINOProjectionHead, | ||
| DINOv2ProjectionHead, | ||
| LeJEPAProjectionHead, | ||
| MIMHead, |
| ) | ||
|
|
||
| criterion = MaskedImageModelingLoss() | ||
| optimizer = torch.optim.AdamW( |
There was a problem hiding this comment.
Since you mention in the comment at the top that the parameters should not require grad, I would apply that here as well. I am aware that the tokenizer uses the @no_grad decorator, but better to completely deactivate the gradients. You can also do it above in the BEIT class.
| from lightly.transforms import BEITTransform | ||
|
|
||
|
|
||
| class BEIT(nn.Module): |
There was a problem hiding this comment.
Would rename to BEiT in order to match the paper's name better.
2 participants
Description
This PR adds support for BEiT (BERT Pre-Training of Image Transformers) for masked image modeling to Lightly.
What was added
Relation to the existing MAE implementation
While both MAE and BEiT are masked image modeling methods, they differ in their reconstruction targets:
References
BEiT: BERT Pre-Training of Image Transformers
https://arxiv.org/abs/2106.08254
Original implementation
https://github.com/microsoft/unilm/tree/master/beit
Documentation
Tests