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,9 +36,13 @@ crossterm = { workspace = true, optional = true }
 # Utilities
 derive-new = { workspace = true }
 serde = { workspace = true, features = ["std", "derive"] }
+strum = { workspace = true }
+yare = "3.0.0"
+
 
 [dev-dependencies]
 burn-ndarray = { path = "../burn-ndarray", version = "0.15.0" }
+approx = "0.5.1"
 
 [package.metadata.docs.rs]
 features = ["doc"]
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::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/base.rs b/crates/burn-train/src/metric/base.rs
@@ -1,4 +1,7 @@
+use burn_core::prelude::{Backend, Bool, Tensor};
+use burn_core::tensor::cast::ToElement;
 use burn_core::{data::dataloader::Progress, LearningRate};
+use strum::EnumIter;
 
 /// Metric metadata that can be used when computing metrics.
 pub struct MetricMetadata {
@@ -56,6 +59,77 @@ pub trait Metric: Send + Sync {
     fn clear(&mut self);
 }
 
+/// The [classification metric](ClassificationMetric) input type.
+#[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(EnumIter, Copy, Clone, Debug)]
+pub enum ClassAverageType {
+    ///Computes the statistics over all classes before averaging
+    Micro,
+    ///Computes the statistics independently for each class before averaging
+    Macro,
+}
+
+impl ClassAverageType {
+    /// sum over samples
+    pub fn aggregate_sum<B: Backend>(self, sample_class_mask: Tensor<B, 2, Bool>) -> Tensor<B, 1> {
+        use ClassAverageType::*;
+        match self {
+            Macro => sample_class_mask.float().sum_dim(0).squeeze(0),
+            Micro => sample_class_mask.float().sum(),
+        }
+    }
+
+    /// average over samples
+    pub fn aggregate_mean<B: Backend>(self, sample_class_mask: Tensor<B, 2, Bool>) -> Tensor<B, 1> {
+        use ClassAverageType::*;
+        match self {
+            Macro => sample_class_mask.float().mean_dim(0).squeeze(0),
+            Micro => sample_class_mask.float().mean(),
+        }
+    }
+
+    ///convert to averaged metric, returns tensor
+    pub fn to_averaged_tensor<B: Backend>(
+        self,
+        mut aggregated_metric: Tensor<B, 1>,
+    ) -> Tensor<B, 1> {
+        use ClassAverageType::*;
+        match self {
+            Macro => {
+                if aggregated_metric.contains_nan().any().into_scalar() {
+                    let nan_mask = aggregated_metric.is_nan();
+                    aggregated_metric = aggregated_metric
+                        .clone()
+                        .select(0, nan_mask.bool_not().argwhere().squeeze(1))
+                }
+                aggregated_metric.mean()
+            }
+            Micro => aggregated_metric,
+        }
+    }
+
+    ///convert to averaged metric, returns float
+    pub fn to_averaged_metric<B: Backend>(self, aggregated_metric: Tensor<B, 1>) -> f64 {
+        self.to_averaged_tensor(aggregated_metric)
+            .into_scalar()
+            .to_f64()
+    }
+}
+
 /// Adaptor are used to transform types so that they can be used by metrics.
 ///
 /// This should be implemented by a model's output type for all [metric inputs](Metric::Input) that are