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-train/Cargo.toml b/crates/burn-train/Cargo.toml
@@ -36,6 +36,9 @@ crossterm = { workspace = true, optional = true }
 # Utilities
 derive-new = { workspace = true }
 serde = { workspace = true, features = ["std", "derive"] }
+rstest.workspace = true
+
+
 
 [dev-dependencies]
 burn-ndarray = { path = "../burn-ndarray", version = "0.15.0" }

diff --git a/crates/burn-train/src/lib.rs b/crates/burn-train/src/lib.rs
@@ -25,3 +25,77 @@ pub use learner::*;
 
 #[cfg(test)]
 pub(crate) type TestBackend = burn_ndarray::NdArray<f32>;
+
+#[cfg(test)]
+pub(crate) type TestDevice = burn_ndarray::NdArrayDevice;
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use crate::{metric::classification::ClassificationInput, TestBackend, TestDevice};
+    use burn_core::{
+        prelude::{Bool, Tensor},
+        tensor::Distribution,
+    };
+    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 (real_targets, prediction_targets) = match classification_type {
+            ClassificationType::Binary => {
+                let real_targets = Tensor::<TestBackend, 2, Bool>::from_data(
+                    [[0], [1], [0], [0], [1]],
+                    &TestDevice::default(),
+                );
+
+                let prediction_targets = Tensor::<TestBackend, 2>::from_data(
+                    [[0], [0], [1], [0], [1]],
+                    &TestDevice::default(),
+                );
+                (real_targets, prediction_targets)
+            }
+            ClassificationType::Multiclass => {
+                let real_targets = Tensor::<TestBackend, 2, Bool>::from_data(
+                    [[0, 1, 0], [1, 0, 0], [0, 0, 1], [0, 0, 1], [1, 0, 0]],
+                    &TestDevice::default(),
+                );
+
+                let prediction_targets = Tensor::<TestBackend, 2>::from_data(
+                    [[0, 1, 0], [0, 1, 0], [1, 0, 0], [0, 0, 1], [1, 0, 0]],
+                    &TestDevice::default(),
+                );
+                (real_targets, prediction_targets)
+            }
+            ClassificationType::Multilabel => {
+                let real_targets = Tensor::<TestBackend, 2, Bool>::from_data(
+                    [[1, 1, 0], [1, 0, 1], [1, 1, 1], [0, 0, 1], [1, 0, 0]],
+                    &TestDevice::default(),
+                );
+
+                let prediction_targets = Tensor::<TestBackend, 2>::from_data(
+                    [[0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 1], [1, 0, 0]],
+                    &TestDevice::default(),
+                );
+                (real_targets, prediction_targets)
+            }
+        };
+        let predictions = prediction_targets
+            .random_like(Distribution::Uniform(0.0, THRESHOLD - 0.1))
+            .sub(prediction_targets.clone())
+            .abs();
+
+        ClassificationInput::new(predictions, real_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(val: ClassificationInput<B>) -> Self {
+        (val.predictions, val.targets)
+    }
+}
+
+/// Class Averaging types for Classification metrics.
+#[derive(Copy, Clone)]
+#[allow(dead_code)]
+pub enum ClassAverageType {
+    ///Computes the statistics over all classes before averaging
+    Micro,
+    ///Computes the statistics independently for each class before averaging
+    Macro,
+}