Tips for those who want to get the last_hidden_state

[MilkClouds]

tldr:

import open_clip
if __name__ == "__main__":
    model, _, preprocess = open_clip.create_model_and_transforms("ViT-L-14", pretrained="datacomp_xl_s13b_b90k")
    model.to("cuda")
    model.eval()

    dummy_input = torch.randn(1, 3, 224, 224, device="cuda")

    visual = model.visual
    visual.output_tokens = True
    visual.attn_pool = visual.proj = None

    pooler_output, last_hidden_state = visual(dummy_input)
    # attach pooler_output to last_hidden_state
    last_hidden_state = torch.cat([pooler_output.unsqueeze(1), last_hidden_state], dim=1)
    print(last_hidden_state.shape) # torch.Size([1, 257, 1024])

explain:
some model(except for vision model for timm and openai) have output_tokens argument. by editing config or simply setting attribute model.output_tokens to True, you can get last_hidden_state(except CLS token). By setting attn_pool and proj to None, we can obtain the CLS token. If you find any errors in this source code, please report me.

[rwightman]

@MilkClouds see also, #721

this is good topic for Discussions, but should probably formalize these into documentation... and make process a bit cleaner (ie support timm properly, although as in 721 you can call a timm specific method on trunk explicitly to get this)