Support returning hidden state in CLIPTextEncoder #442
Conversation
facebook-github-bot
added
the
CLA Signed
Codecov ReportPatch coverage:
Additional details and impacted files@@ Coverage Diff @@
## main #442 +/- ##
==========================================
+ Coverage 68.72% 68.77% +0.05%
==========================================
Files 169 169
Lines 11374 11393 +19
==========================================
+ Hits 7817 7836 +19
Misses 3557 3557
☔ View full report in Codecov by Sentry. |
| ) | ||
| # embeddings now has size [bs, embedding_dim] | ||
|
|
||
| if self.return_hidden_state: | ||
| return embeddings, hidden_state |
| assert isinstance(text_encoder, torch.nn.Module) | ||
|
|
||
| actual_clip_init, actual_hidden_state = text_encoder(text) | ||
| print(actual_hidden_state) |
| expected_hidden_state = torch.Tensor( | ||
| [ | ||
| [ | ||
| [6.348165e-01, -4.137459e-02, -1.604239e00, 1.010798e00], |
There was a problem hiding this comment.
nit: can we limit to 4 decimal places given atol is being set
| @@ -77,6 +82,8 @@ def __init__( | |||
| if use_clip_init: | |||
| self.initialize_parameters() | |||
|
|
|||
| self.return_hidden_state = return_hidden_state | |||
There was a problem hiding this comment.
you can make this an arg to fwd
There was a problem hiding this comment.
+1 to what Ankita said!. We don't need this to be a module level property.
| @@ -120,11 +127,13 @@ def forward(self, text: Tensor) -> Tensor: | |||
|
|
|||
| # [n_ctx, bs, transformer.width] -> [bs, n_ctx, transformer.width] | |||
| embeddings = torch.permute(embeddings, (1, 0, 2)) | |||
| embeddings = self.ln_final(embeddings) | |||
| hidden_state = self.ln_final(embeddings) | |||
| # take features from the eot embedding (the highest number in each sequence) | |||
| embeddings = self.projection( | |||
There was a problem hiding this comment.
I would suggest renaming this to something like projected_embeddings to avoid re-using the embeddings name as it was used before.
|
Thanks both for the review! I've addressed all comments. buck test looks good.
|
|
can you import the diff and land it there |
|
@abhinavarora has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@abhinavarora merged this pull request in a1cc8f3. |
Summary:
Added a
return_hidden_statearg in CLIPTextEncoder. If set to True, forward will return a tuple of final embedding and the last hidden state.Test plan:
pytest tests/models/clip/test_text_encoder.py
cc @abhinavarora Moved code to here for easier review and testing.