burn_train/learner/
classification.rs

1use crate::metric::{
2    AccuracyInput, Adaptor, ConfusionStatsInput, HammingScoreInput, LossInput, PerplexityInput,
3    TopKAccuracyInput, processor::ItemLazy,
4};
5use burn_core::tensor::{Int, Tensor};
6use burn_std::ExecutionError;
7
8/// Simple classification output adapted for multiple metrics.
9///
10/// Supported metrics:
11/// - Accuracy
12/// - AUROC
13/// - TopKAccuracy
14/// - Perplexity
15/// - Precision (via ConfusionStatsInput)
16/// - Recall (via ConfusionStatsInput)
17/// - FBetaScore (via ConfusionStatsInput)
18/// - Loss.
19#[derive(new)]
20pub struct ClassificationOutput {
21    /// The loss.
22    pub loss: Tensor<1>,
23
24    /// The class logits or probabilities. Shape: \[batch_size, num_classes\].
25    pub output: Tensor<2>,
26
27    /// The ground truth class index for each sample. Shape: \[batch_size\].
28    pub targets: Tensor<1, Int>,
29}
30
31impl ItemLazy for ClassificationOutput {
32    fn sync(self) -> Result<Self, ExecutionError> {
33        // No readback: the metrics compute on the device the tensors live on
34        // and read back only their final scalars. Flushing dispatches the
35        // producing stream's buffered work so the metric thread doesn't wait
36        // on an idle queue; all tensors in a training item come off the
37        // autodiff backend entirely, so the metric thread neither retains the
38        // tape nor carries its dispatch context.
39        self.loss.device().flush()?;
40
41        Ok(ClassificationOutput {
42            output: self.output.without_autodiff(),
43            loss: self.loss.without_autodiff(),
44            targets: self.targets.without_autodiff(),
45        })
46    }
47}
48
49impl Adaptor<AccuracyInput> for ClassificationOutput {
50    fn adapt(&self) -> AccuracyInput {
51        AccuracyInput::new(self.output.clone(), self.targets.clone())
52    }
53}
54
55impl Adaptor<LossInput> for ClassificationOutput {
56    fn adapt(&self) -> LossInput {
57        LossInput::new(self.loss.clone())
58    }
59}
60
61impl Adaptor<TopKAccuracyInput> for ClassificationOutput {
62    fn adapt(&self) -> TopKAccuracyInput {
63        TopKAccuracyInput::new(self.output.clone(), self.targets.clone())
64    }
65}
66
67impl Adaptor<PerplexityInput> for ClassificationOutput {
68    fn adapt(&self) -> PerplexityInput {
69        PerplexityInput::new(self.output.clone(), self.targets.clone())
70    }
71}
72
73impl Adaptor<ConfusionStatsInput> for ClassificationOutput {
74    fn adapt(&self) -> ConfusionStatsInput {
75        let [_, num_classes] = self.output.dims();
76        if num_classes > 1 {
77            ConfusionStatsInput::new(
78                self.output.clone(),
79                self.targets.clone().one_hot(num_classes).bool(),
80            )
81        } else {
82            ConfusionStatsInput::new(
83                self.output.clone(),
84                self.targets.clone().unsqueeze_dim(1).bool(),
85            )
86        }
87    }
88}
89
90/// Multi-label classification output adapted for multiple metrics.
91///
92/// Supported metrics:
93/// - HammingScore
94/// - Precision (via ConfusionStatsInput)
95/// - Recall (via ConfusionStatsInput)
96/// - FBetaScore (via ConfusionStatsInput)
97/// - Loss
98#[derive(new)]
99pub struct MultiLabelClassificationOutput {
100    /// The loss.
101    pub loss: Tensor<1>,
102
103    /// The label logits or probabilities. Shape: \[batch_size, num_classes\].
104    pub output: Tensor<2>,
105
106    /// The ground truth labels. Shape: \[batch_size, num_classes\].
107    pub targets: Tensor<2, Int>,
108}
109
110impl ItemLazy for MultiLabelClassificationOutput {
111    fn sync(self) -> Result<Self, ExecutionError> {
112        // Same contract as `ClassificationOutput::sync`: flush and take every
113        // tensor off the autodiff backend, no readback.
114        self.loss.device().flush()?;
115
116        Ok(MultiLabelClassificationOutput {
117            output: self.output.without_autodiff(),
118            loss: self.loss.without_autodiff(),
119            targets: self.targets.without_autodiff(),
120        })
121    }
122}
123
124impl Adaptor<HammingScoreInput> for MultiLabelClassificationOutput {
125    fn adapt(&self) -> HammingScoreInput {
126        HammingScoreInput::new(self.output.clone(), self.targets.clone())
127    }
128}
129
130impl Adaptor<LossInput> for MultiLabelClassificationOutput {
131    fn adapt(&self) -> LossInput {
132        LossInput::new(self.loss.clone())
133    }
134}
135
136impl Adaptor<ConfusionStatsInput> for MultiLabelClassificationOutput {
137    fn adapt(&self) -> ConfusionStatsInput {
138        ConfusionStatsInput::new(self.output.clone(), self.targets.clone().bool())
139    }
140}
burn_train/learner/classification.rs

burn_train/learner/
classification.rs
