Add ViT NEPA#42978
Conversation
|
Hey @yonigozlan that was quick, the PR is ready for a review. This consists mostly in wrapping ViT with the NEPA self supervised framework, pretty straightforward. Test suite is mostly the same as ViT, with an added test for the |
|
cc @molbap as well if you have time to review! |
molbap
left a comment
There was a problem hiding this comment.
Hey @sbucaille ! Thanks for the addition, left a few questions to clarify, if you can help me understand a bit!
| # Mask heads if we want to | ||
| if attention_mask is not None: | ||
| attn_weights = attn_weights * attention_mask |
There was a problem hiding this comment.
sorry, came back to this. I'm a bit confused at the presence of both causal_mask and attention_mask here, especially since at this stage, attention mask has to be added to the weights pre-softmax, right?
There was a problem hiding this comment.
I believe the attention_mask is the head mask which would be a leftover from pre-v5 since head mask support has been dropped if I'm correct. causal_mask in this case is the for the model to attend previous patchs to predict the next (as per the goal of the paper). Since head masks are dropped, attention_mask is always None. But you are right, this is from the original code and the position of these instructions should be above, I fixed it in 835d204 let me know what you think
| query_layer = self.q_norm(query_layer) | ||
| key_layer = self.k_norm(key_layer) |
There was a problem hiding this comment.
this runs regardless of the config.qk_norm value, right? right now it exists and defaults to False, it should inform of the normalization rather (or not exist)
| def forward( | ||
| self, | ||
| pixel_values: torch.Tensor, | ||
| bool_masked_pos: Optional[torch.BoolTensor] = None, | ||
| ) -> torch.Tensor: | ||
| batch_size, _, height, width = pixel_values.shape | ||
| embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=False) | ||
| embeddings_clean = embeddings | ||
|
|
||
| if bool_masked_pos is not None: | ||
| seq_length = embeddings.shape[1] | ||
| mask_tokens = self.mask_token.expand(batch_size, seq_length, -1) | ||
| mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens) | ||
| embeddings = embeddings * (1.0 - mask) + mask_tokens * mask | ||
|
|
||
| # add the [CLS] token | ||
| cls_tokens = self.cls_token.expand(batch_size, -1, -1) | ||
| embeddings = torch.cat((cls_tokens, embeddings), dim=1) | ||
| embeddings_clean = torch.cat((cls_tokens, embeddings_clean), dim=1) | ||
|
|
||
| embeddings = self.dropout(embeddings) | ||
|
|
||
| return embeddings, embeddings_clean |
There was a problem hiding this comment.
this returns a tuple of Tensors (so need to change the hinting). Also returning tuples in module forwards is sometimes causing issues with tracing modules, if both are needed (for training) it would be worth having something like
@dataclass
class ViTNepaEmbeddingOutput:
embeddings: torch.Tensor
clean_embeddings: torch.TensorThere was a problem hiding this comment.
Was not aware of such tracing problems thanks for the heads up, I added the dataclass in 77e4268
| key: torch.Tensor, | ||
| value: torch.Tensor, | ||
| attention_mask: Optional[torch.Tensor], | ||
| scaling: float, |
There was a problem hiding this comment.
if scaling is None, should fallback (like in vit eager)
| if shift: | ||
| # shift one step forward | ||
| p = hidden_states_out[:, :-1, :] # predict next | ||
| z = hidden_states_in[:, 1:, :] # target is next hidden state | ||
| else: | ||
| # same-position matching | ||
| p = hidden_states_out | ||
| z = hidden_states_in |
There was a problem hiding this comment.
do we need the same-position matching? I'm not 100% familiar with the NEPA training setup but we don't need MAE no? (feel free to correct me 😆 )
There was a problem hiding this comment.
You are right, the same-position branch should not exist, shift is never different from True, I removed it in c3279f6
|
Hey @molbap , thanks for your comments ! I addressed them, let me know if it looks good to you ! |
|
Hey @molbap , just bumping this PR as someone asked if it could be merged soon in the original issue |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, vit_nepa |
What does this PR do?
Fixes #42977
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
@yonigozlan