File size: 857 Bytes
6a176cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | import torch.nn.functional as F
class DistillationLoss:
def __init__(self, weight=1.0):
self.weight = weight
def __call__(self, student, teacher):
return self.weight * self._smooth_l1(student, self._detach(teacher))
def _smooth_l1(self, student, teacher):
if isinstance(student, (list, tuple)):
losses = [
F.smooth_l1_loss(student_item, teacher_item)
for student_item, teacher_item in zip(student, teacher)
]
return sum(losses) / len(losses)
return F.smooth_l1_loss(student, teacher)
def _detach(self, teacher):
if isinstance(teacher, list):
return [item.detach() for item in teacher]
if isinstance(teacher, tuple):
return tuple(item.detach() for item in teacher)
return teacher.detach()
|