how to extract features from an image?

lukemelas / PyTorch-Pretrained-ViT

Vision Transformer (ViT) in PyTorch

770 stars 124 forks source link

Closed muaz1994 closed 3 years ago

muaz1994 commented 3 years ago

Hi. I tried to remove the last layer because I only want the features, but I get the error:

The size of tensor a (24) must match the size of tensor b (768) at non-singleton dimension 3

modules = list(model.children())[:-1]
new_model = nn.Sequential(*modules)

Thanks

lukemelas commented 3 years ago

To remove the features, I would recommend just doing del model.fc.