2

我正在尝试使用神经压缩器（intel LPOT）来减小我在 pytorch 中实现的 CNN 模型的大小。我打算做蒸馏

以下是用于提取模型的代码。

    from neural_compressor.experimental import Distillation, common
    from neural_compressor.experimental.common.criterion import PyTorchKnowledgeDistillationLoss
    distiller = Distillation(args.config)
    distiller.student_model = model
    distiller.teacher_model = teacher
    distiller.criterion = PyTorchKnowledgeDistillationLoss()
    distiller.train_func = train_func
    model = distiller.fit()

我想将损失函数更改为不同的损失函数，即我需要提供一个我在 pytorch 中实现的自定义损失函数。目前我在神经压缩器中看到我可以通过向 distiller.criterion 提供参数来改变教师和学生的损失函数，即

    distiller.criterion = PyTorchKnowledgeDistillationLoss(loss_types=['CE', 'KL'])

我认为这是可行的，因为 KullbackLeiblerDivergence 和交叉熵损失在神经压缩器中可用，有什么方法可以提供我的自定义损失函数distiller.criterion吗？

4

1 回答 1

0

在神经压缩器源中有一个名为 PyTorchKnowledgeDistillationLoss 的类，它具有 SoftCrossEntropy 和 KullbackLeiblerDivergence 作为成员函数，如果你想给自己的自定义损失函数添加一个新的成员函数到 PyTorchKnowledgeDistillationLoss 类，它接受 togits 和 target 作为参数，

例如

class PyTorchKnowledgeDistillationLoss(KnowledgeDistillationLoss):
...
...
    def customLossFunction(self, logits, targets):
        //calculate the custom loss
        return custom_loss

然后 PyTorchKnowledgeDistillationLoss 的 init 函数（构造函数）赋值

self.teacher_student_loss = self.customLossFunction
self.student_targets_loss= self.customLossFunction
于 2022-01-24T13:30:23.657 回答