87 lines
3.3 KiB
C++
87 lines
3.3 KiB
C++
#include <sstream>
|
|
#include <algorithm>
|
|
#include <numeric>
|
|
#include <utility>
|
|
#include "RocAuc.h"
|
|
namespace platform {
|
|
std::vector<int> tensorToVector(const torch::Tensor& tensor)
|
|
{
|
|
// Ensure the tensor is of type kInt32
|
|
if (tensor.dtype() != torch::kInt32) {
|
|
throw std::runtime_error("Tensor must be of type kInt32");
|
|
}
|
|
|
|
// Ensure the tensor is contiguous
|
|
torch::Tensor contig_tensor = tensor.contiguous();
|
|
|
|
// Get the number of elements in the tensor
|
|
auto num_elements = contig_tensor.numel();
|
|
|
|
// Get a pointer to the tensor data
|
|
const int32_t* tensor_data = contig_tensor.data_ptr<int32_t>();
|
|
|
|
// Create a std::vector<int> and copy the data
|
|
std::vector<int> result(tensor_data, tensor_data + num_elements);
|
|
|
|
return result;
|
|
}
|
|
double RocAuc::compute(const torch::Tensor& y_proba, const torch::Tensor& labels)
|
|
{
|
|
size_t nClasses = y_proba.size(1);
|
|
// In binary classification problem there's no need to calculate the average of the AUCs
|
|
if (nClasses == 2)
|
|
nClasses = 1;
|
|
size_t nSamples = y_proba.size(0);
|
|
y_test = tensorToVector(labels);
|
|
std::vector<double> aucScores(nClasses, 0.0);
|
|
for (size_t classIdx = 0; classIdx < nClasses; ++classIdx) {
|
|
scoresAndLabels.clear();
|
|
for (size_t i = 0; i < nSamples; ++i) {
|
|
scoresAndLabels.emplace_back(y_proba[i][classIdx].item<float>(), y_test[i] == classIdx ? 1 : 0);
|
|
}
|
|
aucScores[classIdx] = compute_common(nSamples, classIdx);
|
|
}
|
|
return std::accumulate(aucScores.begin(), aucScores.end(), 0.0) / nClasses;
|
|
}
|
|
double RocAuc::compute(const std::vector<std::vector<double>>& y_proba, const std::vector<int>& labels)
|
|
{
|
|
y_test = labels;
|
|
size_t nClasses = y_proba[0].size();
|
|
// In binary classification problem there's no need to calculate the average of the AUCs
|
|
if (nClasses == 2)
|
|
nClasses = 1;
|
|
size_t nSamples = y_proba.size();
|
|
std::vector<double> aucScores(nClasses, 0.0);
|
|
for (size_t classIdx = 0; classIdx < nClasses; ++classIdx) {
|
|
scoresAndLabels.clear();
|
|
for (size_t i = 0; i < nSamples; ++i) {
|
|
scoresAndLabels.emplace_back(y_proba[i][classIdx], labels[i] == classIdx ? 1 : 0);
|
|
}
|
|
aucScores[classIdx] = compute_common(nSamples, classIdx);
|
|
}
|
|
return std::accumulate(aucScores.begin(), aucScores.end(), 0.0) / nClasses;
|
|
}
|
|
double RocAuc::compute_common(size_t nSamples, size_t classIdx)
|
|
{
|
|
std::sort(scoresAndLabels.begin(), scoresAndLabels.end(), std::greater<>());
|
|
std::vector<double> tpr, fpr;
|
|
double tp = 0, fp = 0;
|
|
double totalPos = std::count(y_test.begin(), y_test.end(), classIdx);
|
|
double totalNeg = nSamples - totalPos;
|
|
|
|
for (const auto& [score, label] : scoresAndLabels) {
|
|
if (label == 1) {
|
|
tp += 1;
|
|
} else {
|
|
fp += 1;
|
|
}
|
|
tpr.push_back(tp / totalPos);
|
|
fpr.push_back(fp / totalNeg);
|
|
}
|
|
double auc = 0.0;
|
|
for (size_t i = 1; i < tpr.size(); ++i) {
|
|
auc += 0.5 * (fpr[i] - fpr[i - 1]) * (tpr[i] + tpr[i - 1]);
|
|
}
|
|
return auc;
|
|
}
|
|
} |