Extraction of class token

[vedantdalimkar]

For vision transformer models that have a class token option enabled, how can I extract the class token for intermediate layers?

For example, if I use the model - vit_base_patch14_reg4_dinov2 on a image of height and width of 518, the intermediate feature shape is of B,768,518/14,518/14. Where is the class token in this feature?

[rwightman]

@vedantdalimkar assuming you're using forward_intermediates(), return_prefix_tokens=True will give you another list of the class tokens beside the spatial tokens

pytorch-image-models/timm/models/vision_transformer.py

Lines 722 to 743 in e44f14d

	def forward_intermediates(
	self,
	x: torch.Tensor,
	indices: Optional[Union[int, List[int]]] = None,
	return_prefix_tokens: bool = False,
	norm: bool = False,
	stop_early: bool = False,
	output_fmt: str = 'NCHW',
	intermediates_only: bool = False,
	) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
	""" Forward features that returns intermediates.
	Args:
	x: Input image tensor
	indices: Take last n blocks if int, all if None, select matching indices if sequence
	return_prefix_tokens: Return both prefix and spatial intermediate tokens
	norm: Apply norm layer to all intermediates
	stop_early: Stop iterating over blocks when last desired intermediate hit
	output_fmt: Shape of intermediate feature outputs
	intermediates_only: Only return intermediate features
	Returns:

[vedantdalimkar]

Is there a model config setup that will give the class token through forward function?

[rwightman]

@vedantdalimkar no, I didn't want to risk breaking backwards compat or causing typing issues on the main forward, so need to use forward_intermediates, forward intermediates can chain with forward_head so you can get the combo of intermediates + the final classification result in a train loop if you want.

And if you do want intermediate tokens available via model() it's best to use a small wrapper nn.Module that calls forward_intermediates and returns just the class tokens, or whatever else you want in whatever form you prefer.

[rwightman]

@vedantdalimkar oh, and there are ways to get the features w/o using intermediates, but only the last hidden state ... using forward_features or normal forward but with pooling and no classifier... it returns the flat sequence so class token is [0] and the 4 register tokens are [1:5].

mm = timm.create_model('vit_base_patch14_reg4_dinov2', pretrained=True)
oo = mm.forward_features(torch.randn(2,3,518,518))
mm2 = timm.create_model('vit_base_patch14_reg4_dinov2', pretrained=True, global_pool='', num_classes=0)
oo2 = mm2(torch.randn(2,3,518,518))
oo.shape
torch.Size([2, 1374, 768])
oo2.shape
torch.Size([2, 1374, 768])