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
@@ -7,6 +7,7 @@ throughout the training process. We currently offer a restricted range of metric
 | ---------------- | ------------------------------------------------------- |
 | Accuracy         | Calculate the accuracy in percentage                    |
 | TopKAccuracy     | Calculate the top-k accuracy in percentage              |
+| Precision        | Calculate precision in percentage                       |
 | AUROC            | Calculate the area under curve of ROC in percentage     |
 | Loss             | Output the loss used for the backward pass              |
 | CPU Temperature  | Fetch the temperature of CPUs                           |

diff --git a/crates/burn-tensor/src/tensor/api/check.rs b/crates/burn-tensor/src/tensor/api/check.rs
@@ -478,6 +478,11 @@ impl TensorCheck {
                     "Can't create a one hot tensor from ({index_tensor:?}) containing indexes greater or equal to the number of classes ({num_classes})",
                 )),
             );
+        } else if num_classes <= 1 {
+            check = check.register(
+                "One Hot",
+                TensorError::new("Can't create a one hot tensor with less then 2 classes"),
+            )
         }
         check
     }

diff --git a/crates/burn-train/Cargo.toml b/crates/burn-train/Cargo.toml
@@ -40,6 +40,7 @@ ratatui = { workspace = true, optional = true, features = ["all-widgets", "cross
 derive-new = { workspace = true }
 serde = { workspace = true, features = ["std", "derive"] }
 async-channel = { workspace = true }
+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,4 +1,4 @@
-use crate::metric::{AccuracyInput, Adaptor, HammingScoreInput, LossInput};
+use crate::metric::{AccuracyInput, Adaptor, HammingScoreInput, LossInput, PrecisionInput};
 use burn_core::tensor::backend::Backend;
 use burn_core::tensor::{Int, Tensor};
 
@@ -27,6 +27,23 @@ impl<B: Backend> Adaptor<LossInput<B>> for ClassificationOutput<B> {
     }
 }
 
+impl<B: Backend> Adaptor<PrecisionInput<B>> for ClassificationOutput<B> {
+    fn adapt(&self) -> PrecisionInput<B> {
+        let [_, num_classes] = self.output.dims();
+        if num_classes > 1 {
+            PrecisionInput::new(
+                self.output.clone(),
+                self.targets.clone().one_hot(num_classes).bool(),
+            )
+        } else {
+            PrecisionInput::new(
+                self.output.clone(),
+                self.targets.clone().unsqueeze_dim(1).bool(),
+            )
+        }
+    }
+}
+
 /// Multi-label classification output adapted for multiple metrics.
 #[derive(new)]
 pub struct MultiLabelClassificationOutput<B: Backend> {
@@ -51,3 +68,9 @@ impl<B: Backend> Adaptor<LossInput<B>> for MultiLabelClassificationOutput<B> {
         LossInput::new(self.loss.clone())
     }
 }
+
+impl<B: Backend> Adaptor<PrecisionInput<B>> for MultiLabelClassificationOutput<B> {
+    fn adapt(&self) -> PrecisionInput<B> {
+        PrecisionInput::new(self.output.clone(), self.targets.clone().bool())
+    }
+}
diff --git a/crates/burn-train/src/lib.rs b/crates/burn-train/src/lib.rs
@@ -26,3 +26,78 @@ pub use learner::*;
 
 #[cfg(test)]
 pub(crate) type TestBackend = burn_ndarray::NdArray<f32>;
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use crate::TestBackend;
+    use burn_core::{prelude::Tensor, tensor::Bool};
+    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,
+    ) -> (Tensor<TestBackend, 2>, Tensor<TestBackend, 2, Bool>) {
+        match classification_type {
+            ClassificationType::Binary => {
+                (
+                    Tensor::from_data(
+                        [[0.3], [0.2], [0.7], [0.1], [0.55]],
+                        //[[0],   [0],   [1],   [0],   [1]] with threshold=0.5
+                        &Default::default(),
+                    ),
+                    Tensor::from_data([[0], [1], [0], [0], [1]], &Default::default()),
+                )
+            }
+            ClassificationType::Multiclass => {
+                (
+                    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(),
+                    ),
+                    Tensor::from_data(
+                        [[0, 1, 0], [1, 0, 0], [0, 0, 1], [0, 0, 1], [1, 0, 0]],
+                        &Default::default(),
+                    ),
+                )
+            }
+            ClassificationType::Multilabel => {
+                (
+                    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(),
+                    ),
+                    Tensor::from_data(
+                        [[1, 1, 0], [1, 0, 1], [1, 1, 1], [0, 0, 1], [1, 0, 0]],
+                        &Default::default(),
+                    ),
+                )
+            }
+        }
+    }
+}
diff --git a/crates/burn-train/src/metric/classification.rs b/crates/burn-train/src/metric/classification.rs
@@ -0,0 +1,9 @@
+/// The reduction strategy for classification metrics.
+#[derive(Copy, Clone, Default)]
+pub enum ClassReduction {
+    /// Computes the statistics over all classes before averaging
+    Micro,
+    /// Computes the statistics independently for each class before averaging
+    #[default]
+    Macro,
+}