Add another MLP head in vision transformer

lucidrains / vit-pytorch

Implementation of Vision Transformer, a simple way to achieve SOTA in vision classification with only a single transformer encoder, in Pytorch

MIT License

20.57k stars 3.05k forks source link

Add another MLP head in vision transformer #222

Open Atul997 opened 2 years ago

Atul997 commented 2 years ago

How can I add/extend MLP head in same model for detection? Let's say head is detecting objects A,B,C in a image and we want to train by adding or extend MLP/classification head to detect objects D, E, F in same image without combining both the classes.

Any help would be appreciated?

lucidrains commented 2 years ago

@Atul997 this is a nice scheme https://github.com/lucidrains/vit-pytorch#learnable-memory-vit

Atul997 commented 2 years ago

@lucidrains Thanks . How the data will look like in terms of json file such as two different json or single json containing annotation for both the different head classes?