60 lines
1.9 KiB
Rust
60 lines
1.9 KiB
Rust
/// Multi-label classification output adapted for multiple metrics.
|
|
///
|
|
/// Supported metrics:
|
|
/// - HammingScore
|
|
/// - Precision (via ConfusionStatsInput)
|
|
/// - Recall (via ConfusionStatsInput)
|
|
/// - FBetaScore (via ConfusionStatsInput)
|
|
/// - Loss
|
|
#[derive(new)]
|
|
pub struct MultiLabelSoftClassificationOutput<B: Backend> {
|
|
/// The loss.
|
|
pub loss: Tensor<B, 1>,
|
|
|
|
/// The label logits or probabilities. Shape: \[batch_size, num_classes\].
|
|
pub output: Tensor<B, 2>,
|
|
|
|
/// The ground truth labels (target values). Shape: \[batch_size, num_classes\].
|
|
pub targets: Tensor<B, 2>,
|
|
}
|
|
|
|
impl<B: Backend> ItemLazy for MultiLabelSoftClassificationOutput<B> {
|
|
type ItemSync = MultiLabelSoftClassificationOutput<NdArray>;
|
|
|
|
fn sync(self) -> Self::ItemSync {
|
|
let [output, loss, targets] = Transaction::default()
|
|
.register(self.output)
|
|
.register(self.loss)
|
|
.register(self.targets)
|
|
.execute()
|
|
.try_into()
|
|
.expect("Correct amount of tensor data");
|
|
|
|
let device = &Default::default();
|
|
|
|
MultiLabelSoftClassificationOutput {
|
|
output: Tensor::from_data(output, device),
|
|
loss: Tensor::from_data(loss, device),
|
|
targets: Tensor::from_data(targets, device),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<B: Backend> Adaptor<HammingScoreInput<B>> for MultiLabelSoftClassificationOutput<B> {
|
|
fn adapt(&self) -> HammingScoreInput<B> {
|
|
HammingScoreInput::new(self.output.clone(), self.targets.clone())
|
|
}
|
|
}
|
|
|
|
impl<B: Backend> Adaptor<LossInput<B>> for MultiLabelSoftClassificationOutput<B> {
|
|
fn adapt(&self) -> LossInput<B> {
|
|
LossInput::new(self.loss.clone())
|
|
}
|
|
}
|
|
|
|
impl<B: Backend> Adaptor<ConfusionStatsInput<B>> for MultiLabelSoftClassificationOutput<B> {
|
|
fn adapt(&self) -> ConfusionStatsInput<B> {
|
|
ConfusionStatsInput::new(self.output.clone(), self.targets.clone().bool())
|
|
}
|
|
}
|