search

lightly-ai / lightly

A python library for self-supervised learning on images.
https://docs.lightly.ai/self-supervised-learning/
MIT License
2.89k stars 249 forks source link

random_token_mask from models.utils masks a patch too much #1583

Closed Rob2U closed 1 month ago

Rob2U commented 1 month ago

random_token_mask from models.utils masks a patch too much when mask_class_token=False.

Input: size = (batch_size, sequence_length) <- sequence length includes class token mask_ratio = 0.75 mask_class_token = False

Output shapes: idx_keep.shape: (batch_size, 64) idx_mask.shape: (batch_size, 193)

I would expect to exclude the class_token entirely from masking and from mask calculation. Therefore, the output should have shapes: idx_keep.shape: (batch_size, 65) idx_mask.shape: (batch_size, 192)

Additionally, the return type annotation seems to be outdated (torch.Tensor instead of Tuple[torch.Tensor, torch.Tensor])

Thank you in advance, and especially thank you for this fantastic repo.

Rob2U commented 1 month ago

https://github.com/lightly-ai/lightly/blob/master/lightly/models/utils.py#L464

guarin commented 1 month ago

Hi! Thanks for the issue, created a PR to fix it #1585 :)