Customizable Loss Module

[AhmedHussKhalifa]

Hi,

I am trying to customize a new form of KD loss function. I had been referred to this link as an example, but I still not there yet. I have used to change the locally installed torchditill instead of its PyPI package version. It is little bit changing to me to understand how to implement.

I took KDloss as a reference, where I registered to @register_single_loss.

This is my implementation:

from torchdistill.losses.registry import register_single_loss

@register_single_loss
class KDLoss(nn.KLDivLoss):
    def __init__(self, student_module_path, student_module_io, teacher_module_path, teacher_module_io,
                 temperature, alpha=None, beta=None, reduction='batchmean', **kwargs):
        super().__init__(reduction=reduction)
        self.student_module_path = student_module_path
        self.student_module_io = student_module_io
        self.teacher_module_path = teacher_module_path
        self.teacher_module_io = teacher_module_io
        self.temperature = temperature
        self.alpha = alpha
        self.beta = 1 - alpha if beta is None else beta
        cel_reduction = 'mean' if reduction == 'batchmean' else reduction
        self.cross_entropy_loss = nn.CrossEntropyLoss(reduction=cel_reduction, **kwargs)

    def forward(self, student_io_dict, teacher_io_dict, targets=None, *args, **kwargs):
        student_logits = student_io_dict[self.student_module_path][self.student_module_io]
        teacher_logits = teacher_io_dict[self.teacher_module_path][self.teacher_module_io]
        soft_loss = super().forward(torch.log_softmax(student_logits / self.temperature, dim=1),
                                    torch.softmax(teacher_logits / self.temperature, dim=1))
        if self.alpha is None or self.alpha == 0 or targets is None:
            return soft_loss

        hard_loss = self.cross_entropy_loss(student_logits, targets)
        return self.alpha * hard_loss + self.beta * (self.temperature ** 2) * soft_loss

This is the error message;

  File "./torchdistill/examples/KD_Loss.py", line 2, in <module>
    from torchdistill.losses.registry import register_single_loss
ImportError: cannot import name 'register_single_loss' from 'torchdistill.losses.registry' (/home/ahamsala/PROJECTS/torchdistill_env/lib/python3.8/site-packages/torchdistill/losses/registry.py)

[yoshitomo-matsubara]

Hi @AhmedHussKhalifa

I need the version of torchdistill you are using to answer the question.
Also, as I directly suggested to you here, please close the existing discussions at first (choose answers or respond if the answers were not satisfactory).

If you would not do so, I cannot spend any more of my time for your questions.

[AhmedHussKhalifa]

Hi @yoshitomo-matsubara ,

I'm sorry if I've done something to make things difficult or to bother you. I am using the most recent one. However, I am using an older version for now (0.3.2), but I am still curious about how to do add an extra module, if possible.

[yoshitomo-matsubara]

@AhmedHussKhalifa

Before I answer your questions,

1. Check your discussions
2. Confirm that you have closed only this discussion, but not closed discussions 1 2 3
3. Visit each of the "open" discussions
4. Mark as answer or respond if the answers were not satisfactory

So far, you opened multiple threads in "Discussions" and asked for help. But you didn't mark as answer or respond to the answers either

[AhmedHussKhalifa]

Done

[yoshitomo-matsubara]

If you're using torchdistill<=0.3.3, the first line should be

from torchdistill.losses.single import register_single_loss

See https://github.com/yoshitomo-matsubara/torchdistill/blob/v0.3.3/torchdistill/losses/single.py