temperature_scaling.py 文件源码

def set_temperature(self, valid_loader):
        """
        Tune the tempearature of the model (using the validation set).
        We're going to set it to optimize NLL.
        valid_loader (DataLoader): validation set loader
        """
        self.cuda()
        nll_criterion = nn.CrossEntropyLoss().cuda()
        ece_criterion = _ECELoss().cuda()

        # First: collect all the logits and labels for the validation set
        logits_list = []
        labels_list = []
        for input, label in valid_loader:
            input_var = Variable(input, volatile=True).cuda()
            logits_var = self.model(input_var)
            logits_list.append(logits_var.data)
            labels_list.append(label)
        logits = torch.cat(logits_list).cuda()
        labels = torch.cat(labels_list).cuda()
        logits_var = Variable(logits)
        labels_var = Variable(labels)

        # Calculate NLL and ECE before temperature scaling
        before_temperature_nll = nll_criterion(logits_var, labels_var).data[0]
        before_temperature_ece = ece_criterion(logits_var, labels_var).data[0]
        print('Before temperature - NLL: %.3f, ECE: %.3f' % (before_temperature_nll, before_temperature_ece))

        # Next: optimize the temperature w.r.t. NLL
        optimizer = optim.LBFGS([self.temperature], lr=0.01, max_iter=50)
        def eval():
            loss = nll_criterion(self.temperature_scale(logits_var), labels_var)
            loss.backward()
            return loss
        optimizer.step(eval)

        # Calculate NLL and ECE after temperature scaling
        after_temperature_nll = nll_criterion(self.temperature_scale(logits_var), labels_var).data[0]
        after_temperature_ece = ece_criterion(self.temperature_scale(logits_var), labels_var).data[0]
        print('Optimal temperature: %.3f' % self.temperature.data[0])
        print('After temperature - NLL: %.3f, ECE: %.3f' % (after_temperature_nll, after_temperature_ece))

        return self