How to delete last linear layer of ViT-H or how to replace it with Identity()?

Hello there! I’m trying to figure out how I could possibly replace or delete the last layer of VIT-H or ViT-b.
Here deleting children method doesn’t work like Resnet.


Thanks in advance for the suggestion!

looking for an answer!

You can directly replace layers with an nn.Identity module.
I would also not recommend bumping the thread after 2h of wait time, as it’s just spamming the board and not helpful.

I can access to model.encoder layers. But I didn’t find any solution how to access the last linear layer in VIT!

If you are loading the vision transformer from torchvision, try modifying the linear layer by model.heads = torch.nn.Identity(). If you check the source code (I’m looking at version 0.15.2), you’ll see that the output of model.encoder is passed to model.heads.

In general, you can find out the attributes of a Python object by obj.__dict__.