Could you explain the code of Mask Class Token in the project?

[FYY799]

Could you explain the code of Mask Class Token in the project?

[zh-ding]

There are two main parts of the Mask Class Token. First one is that we add mask class token (by repeating the class token) to CLIP

MaskCLIP/maskclip/modeling/maskclip.py

Lines 165 to 168 in 43049d1

	x = x.permute(1, 0, 2) # NLD -> LND
	cls_embed = x[0:1]
	cls_embed = cls_embed.repeat(q, 1, 1)
	x = torch.cat([cls_embed, x], dim=0)

Second one is that we use the segmentation masks to serve as attention mask in the self-attention layer

MaskCLIP/maskclip/modeling/maskclip.py

Lines 50 to 66 in 43049d1

	def forward(self, x: torch.Tensor, masks: torch.Tensor, masks_embed: torch.Tensor = None):
	l, b, d = x.shape
	_, q, _, _ = masks.shape
	masks = (masks.sigmoid() >= 0.5).float()
	masks = F.max_pool2d(masks, self.clip_patch_size).flatten(2)
	attn_mask = torch.empty((b, l, l), device=x.device, dtype=torch.bool)
	attn_mask[:, :, :] = False
	attn_mask[:, :, :q] = True
	attn_mask[:, :q, q+1:] = masks == 0.
	attn_mask = torch.repeat_interleave(attn_mask, self.n_head, dim=0)
	x_res, masks_res = self.attention(self.ln_1(x), attn_mask=attn_mask, masks_embed=masks_embed)
	x = x + x_res
	x = x + self.mlp(self.ln_2(x))
	return x, masks_res

.