tracel-ai · laggui · Nov 20, 2024 · Aug 21, 2024 · Sep 9, 2024 · Sep 9, 2024
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/burn-book/src/building-blocks/metric.md b/burn-book/src/building-blocks/metric.md
@@ -4,9 +4,10 @@ When working with the learner, you have the option to record metrics that will b
 throughout the training process. We currently offer a restricted range of metrics.
 
 | Metric           | Description                                             |
-| ---------------- | ------------------------------------------------------- |
+|------------------|---------------------------------------------------------|
 | Accuracy         | Calculate the accuracy in percentage                    |
 | TopKAccuracy     | Calculate the top-k accuracy in percentage              |
+| Precision        | Calculate precision in percentage                       |
 | Loss             | Output the loss used for the backward pass              |
 | CPU Temperature  | Fetch the temperature of CPUs                           |
 | CPU Usage        | Fetch the CPU utilization                               |

diff --git a/crates/burn-tensor/src/tensor/api/int.rs b/crates/burn-tensor/src/tensor/api/int.rs
@@ -52,10 +52,14 @@ where
     /// ```
     pub fn one_hot(self, num_classes: usize) -> Tensor<B, 2, Int> {
         check!(TensorCheck::one_hot_tensor(self.clone(), num_classes));
-        let [num_samples] = self.dims();
-        let indices = self.unsqueeze();
-        let values = indices.ones_like();
-        Tensor::zeros([num_samples, num_samples], &indices.device()).scatter(1, indices, values)
+        if num_classes == 1 {
+            self.unsqueeze_dim(1)
+        } else {
+            let [num_samples] = self.dims();
+            let indices = self.unsqueeze_dim(1);
+            let values = indices.ones_like();
+            Tensor::zeros([num_samples, num_classes], &indices.device()).scatter(1, indices, values)
+        }
     }
 }
 

diff --git a/crates/burn-train/Cargo.toml b/crates/burn-train/Cargo.toml
@@ -39,6 +39,9 @@ ratatui = { workspace = true, optional = true, features = ["all-widgets", "cross
 # Utilities
 derive-new = { workspace = true }
 serde = { workspace = true, features = ["std", "derive"] }
+rstest.workspace = true
+
+
 
 [dev-dependencies]
 burn-ndarray = { path = "../burn-ndarray", version = "0.16.0" }

diff --git a/crates/burn-train/src/learner/classification.rs b/crates/burn-train/src/learner/classification.rs
@@ -1,3 +1,4 @@
+use crate::metric::classification::ClassificationInput;
 use crate::metric::{AccuracyInput, Adaptor, HammingScoreInput, LossInput};
 use burn_core::tensor::backend::Backend;
 use burn_core::tensor::{Int, Tensor};
@@ -27,6 +28,16 @@ impl<B: Backend> Adaptor<LossInput<B>> for ClassificationOutput<B> {
     }
 }
 
+impl<B: Backend> Adaptor<ClassificationInput<B>> for ClassificationOutput<B> {
+    fn adapt(&self) -> ClassificationInput<B> {
+        let [_, num_classes] = self.output.dims();
+        ClassificationInput::new(
+            self.output.clone(),
+            self.targets.clone().one_hot(num_classes).bool(),
+        )
+    }
+}
+
 /// Multi-label classification output adapted for multiple metrics.
 #[derive(new)]
 pub struct MultiLabelClassificationOutput<B: Backend> {

diff --git a/crates/burn-train/src/lib.rs b/crates/burn-train/src/lib.rs
@@ -26,3 +26,79 @@ pub use learner::*;
 
 #[cfg(test)]
 pub(crate) type TestBackend = burn_ndarray::NdArray<f32>;
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use crate::{metric::classification::ClassificationInput, TestBackend};
+    use burn_core::prelude::Tensor;
+    use std::default::Default;
+
+    /// Probability of tp before adding errors
+    pub const THRESHOLD: f64 = 0.5;
+
+    #[derive(Debug)]
+    pub enum ClassificationType {
+        Binary,
+        Multiclass,
+        Multilabel,
+    }
+
+    /// Sample x Class shaped matrix for use in
+    /// classification metrics testing
+    pub fn dummy_classification_input(
+        classification_type: &ClassificationType,
+    ) -> ClassificationInput<TestBackend> {
+        let (targets, predictions) = match classification_type {
+            ClassificationType::Binary => {
+                (
+                    Tensor::from_data([[0], [1], [0], [0], [1]], &Default::default()),
+                    Tensor::from_data(
+                        [[0.3], [0.2], [0.7], [0.1], [0.55]],
+                        //[[0],   [0],   [1],   [0],   [1]] with threshold=0.5
+                        &Default::default(),
+                    ),
+                )
+            }
+            ClassificationType::Multiclass => {
+                (
+                    Tensor::from_data(
+                        [[0, 1, 0], [1, 0, 0], [0, 0, 1], [0, 0, 1], [1, 0, 0]],
+                        &Default::default(),
+                    ),
+                    Tensor::from_data(
+                        [
+                            [0.2, 0.8, 0.0],
+                            [0.3, 0.6, 0.1],
+                            [0.7, 0.25, 0.05],
+                            [0.1, 0.15, 0.8],
+                            [0.9, 0.03, 0.07],
+                        ],
+                        //[[0,   1,   0],   [0,   1,   0],    [1,   0,   0],    [0,   0,   1],    [1,   0,    0]] with top_k=1
+                        //[[1,   1,   0],   [1,   1,   0],    [1,   1,   0],    [0,   1,   1],    [1,   0,    1]] with top_k=2
+                        &Default::default(),
+                    ),
+                )
+            }
+            ClassificationType::Multilabel => {
+                (
+                    Tensor::from_data(
+                        [[1, 1, 0], [1, 0, 1], [1, 1, 1], [0, 0, 1], [1, 0, 0]],
+                        &Default::default(),
+                    ),
+                    Tensor::from_data(
+                        [
+                            [0.1, 0.7, 0.6],
+                            [0.3, 0.9, 0.05],
+                            [0.8, 0.9, 0.4],
+                            [0.7, 0.5, 0.9],
+                            [1.0, 0.3, 0.2],
+                        ],
+                        //[[0,   1,   1],   [0,   1,   0],    [1,   1,   0],   [1,   0,   1],   [1,   0,   0]] with threshold=0.5
+                        &Default::default(),
+                    ),
+                )
+            }
+        };
+        ClassificationInput::new(predictions, targets)
+    }
+}
diff --git a/crates/burn-train/src/metric/classification.rs b/crates/burn-train/src/metric/classification.rs
@@ -0,0 +1,26 @@
+use burn_core::prelude::{Backend, Bool, Tensor};
+
+/// Input for classification tasks.
+#[derive(new, Debug, Clone)]
+pub struct ClassificationInput<B: Backend> {
+    /// Sample x Class Non thresholded normalized predictions.
+    pub predictions: Tensor<B, 2>,
+    /// Sample x Class one-hot encoded target.
+    pub targets: Tensor<B, 2, Bool>,
+}
+
+impl<B: Backend> From<ClassificationInput<B>> for (Tensor<B, 2>, Tensor<B, 2, Bool>) {
+    fn from(input: ClassificationInput<B>) -> Self {
+        (input.predictions, input.targets)
+    }
+}
+
+/// Class Averaging types for Classification metrics.
+#[derive(Copy, Clone, Default)]
+pub enum ClassReduction {
+    ///Computes the statistics over all classes before averaging
+    #[default]
+    Micro,
+    ///Computes the statistics independently for each class before averaging
+    Macro,
+}