calibration_loss.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstdint>
#include <vector>
 
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
class CalibrationLoss {
public:
    explicit CalibrationLoss(DeviceBackend* backend);
 
    [[nodiscard]] float forward(const Tensor& probs, const Tensor& target_class);
 
    [[nodiscard]] Tensor backward() const;
 
private:
    DeviceBackend* backend_;
    Tensor last_probs_;
    std::vector<int64_t> last_target_indices_;
    Tensor last_one_hot_ = Tensor(Shape({0}), backend_);  // (N, num_classes), on the probs' device
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix