search

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
19.5k stars 2.95k forks source link

Questions about distill_loss #289

Open haoren55555 opened 9 months ago

haoren55555 commented 9 months ago

sorry to bother, I see the distill_loss in distill.py as : distill_loss = F.kl_div( F.log_softmax(distill_logits / T, dim=-1), F.softmax(teacher_logits / T, dim=-1).detach(), reduction='batchmean') I wonder why the teacher part uses the softmax function rather than log_softmax one, thanks.

lucidrains commented 9 months ago

@haoren55555 yeah, you could do log softmax for teacher too by setting log_target = True https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html . i'm just rolling with what pytorch offers

vivekh2000 commented 2 months ago

That was insightful. However, in your implementation of class DistillMixin, you are passing the cls token from the MLP lead.

The MLP head of the distillation token is implemented using LayerNorm, whereas the cls token MLP head is implemented in vit.py without LayerNorm. Why is it so? What am I missing? Please correct me if I am wrong. image

image

image

vivekh2000 commented 2 months ago

@lucidrains I have also noticed that in the paper, the authors mentioned that at test time, they fused the two heads, i.e., MLP heads for cls token and distilled token.

for which we add the softmax output by the two classifiers to make the prediction. (p.8, paragraph 3)

This fusion of two heads is not implemented in the forward method of the DistillMixin class. Please let me know if it is implemented elsewhere or if I am missing something. Thanks for the excellent codebase.