Getting a fixed size embedding from the last hidden state.
Author: PrithivirajDamodaran
Pull request: no

I am trying to play with the bare vision model's last hidden state (like ViT, BEiT). How can I get a fixed size representation out of the last_hidden_state which is of shape (channels, height, width) say of 1D? Should I simply flatten it? 

Please advice

Comments:
Hi,

First of all, please use the forum for these kind of questions, we'd like to keep Github issues for bugs/feature requests.

Second, a last hidden state is typically of shape (batch_size, seq_len, hidden_size) for these kind of models, as they are Transformer-based. You feed a sequence of patches through a Transformer encoder, hence you end up with a vector for each of these patches at the end. You can permute it to turn it into a tensor of shape (batch_size, hidden_size, seq_len), and split the last dimension based on the patch size, to get a tensor of shape (batch_size, hidden_size, patch_size, patch_size). This gives you an image-like representation.

In code:

```
from transformers import AutoModel
import torch

model = AutoModel.from_pretrained("microsoft/beit-base-patch16-224")

pixel_values = torch.randn(1, 3, 224, 224)

outputs = model(pixel_values)
last_hidden_state = outputs.last_hidden_state[:,1:,:] # we discard the CLS token

batch_size = last_hidden_state.shape[0]
num_patches = model.config.image_size // model.config.patch_size
image_like_representation = last_hidden_state.permute(0, 2, 1)
image_like_representation = image_like_representation.view(batch_size, -1, num_patches, num_patches)
```

Rank Apology - I will use the forum going forward.

But just to close this one out, maybe I am not explaining properly: I am using BEiT as the image encoder and BERT as the Text encoder. I am trying to get a fixed size 1D representation of the last_hidden_state of BEiT (Something like what CLIP obtains) to concatenate with a BERT embedding to feed into an MLP head.

I could use the pooler_output of the image but it doesn’t seem to preserve certain spatial nuances hence I would like to use the last_hidden_state


Can you help me? Moved to the forum as well - https://discuss.huggingface.co/t/how-to-get-a-fixed-size-embedding-from-the-last-hidden-state-of-vision-models/17275


@PrithivirajDamodaran you can adapt code snippet from @NielsRogge as follows:

```
from transformers import AutoModel
import torch

model = AutoModel.from_pretrained("microsoft/beit-base-patch16-224")

pixel_values = torch.randn(1, 3, 224, 224)

outputs = model(pixel_values)
img_embedding = outputs.last_hidden_state[0, 0, :] # CLS token of the last layer can be used as the image embedding

```

@nihit - To get the CLS token, I can directly get the ```pooler_output```, instead of slicing last_hidden_state. Because all HF bare models returns by default two keys ```pooler_output``` and ```last_hidden_state```.  BTW pooler_output is nothing but the raw first entry (CLS) from the last_hidden_state but it was passed through a simple MLP pooler layer (linear + tanh). 

I am NOT interested in CLS embedding, I would like to have a fixed representation of the entire last layer itself i.e. last_hidden_state.

I have figured this out.

@PrithivirajDamodaran - I noticed that sometimes the last_hidden_state has seq length dimension different from the number of input_ids from the pre-processor (for example layoutlm). Does anyone know why this happens?