Clean up train system metrics (#2707)

crates/burn-train/Cargo.toml

@@ -12,9 +12,9 @@ documentation = "https://docs.rs/burn-train"
 version.workspace = true
 
 [features]
-default = ["metrics", "tui"]
+default = ["sys-metrics", "tui"]
 doc = ["default"]
-metrics = ["nvml-wrapper", "sysinfo", "systemstat"]
+sys-metrics = ["nvml-wrapper", "sysinfo", "systemstat"]
 tui = ["ratatui"]
 
 [dependencies]
@@ -28,7 +28,7 @@ tracing-subscriber = { workspace = true }
 tracing-appender = { workspace = true }
 tracing-core = { workspace = true }
 
-# Metrics
+# System Metrics
 nvml-wrapper = { workspace = true, optional = true }
 sysinfo = { workspace = true, optional = true }
 systemstat = { workspace = true, optional = true }

crates/burn-train/src/metric/mod.rs

@@ -3,65 +3,54 @@ pub mod state;
 /// Module responsible to save and exposes data collected during training.
 pub mod store;
 
+// System metrics
+#[cfg(feature = "sys-metrics")]
+mod cpu_temp;
+#[cfg(feature = "sys-metrics")]
+mod cpu_use;
+#[cfg(feature = "sys-metrics")]
+mod cuda;
+#[cfg(feature = "sys-metrics")]
+mod memory_use;
+#[cfg(feature = "sys-metrics")]
+pub use cpu_temp::*;
+#[cfg(feature = "sys-metrics")]
+pub use cpu_use::*;
+#[cfg(feature = "sys-metrics")]
+pub use cuda::*;
+#[cfg(feature = "sys-metrics")]
+pub use memory_use::*;
+
+// Training metrics
 mod acc;
 mod auroc;
 mod base;
-#[cfg(feature = "metrics")]
 mod confusion_stats;
-#[cfg(feature = "metrics")]
-mod cpu_temp;
-#[cfg(feature = "metrics")]
-mod cpu_use;
-#[cfg(feature = "metrics")]
-mod cuda;
-#[cfg(feature = "metrics")]
 mod fbetascore;
 mod hamming;
-#[cfg(feature = "metrics")]
 mod iteration;
 mod learning_rate;
 mod loss;
-#[cfg(feature = "metrics")]
-mod memory_use;
-#[cfg(feature = "metrics")]
 mod precision;
-#[cfg(feature = "metrics")]
 mod recall;
-#[cfg(feature = "metrics")]
 mod top_k_acc;
 
 pub use acc::*;
 pub use auroc::*;
 pub use base::*;
-#[cfg(feature = "metrics")]
 pub use confusion_stats::ConfusionStatsInput;
-#[cfg(feature = "metrics")]
-pub use cpu_temp::*;
-#[cfg(feature = "metrics")]
-pub use cpu_use::*;
-#[cfg(feature = "metrics")]
-pub use cuda::*;
-#[cfg(feature = "metrics")]
 pub use fbetascore::*;
 pub use hamming::*;
-#[cfg(feature = "metrics")]
 pub use iteration::*;
 pub use learning_rate::*;
 pub use loss::*;
-#[cfg(feature = "metrics")]
-pub use memory_use::*;
-#[cfg(feature = "metrics")]
 pub use precision::*;
-#[cfg(feature = "metrics")]
 pub use recall::*;
-#[cfg(feature = "metrics")]
 pub use top_k_acc::*;
 
-#[cfg(feature = "metrics")]
 pub(crate) mod classification;
 pub(crate) mod processor;
 
-#[cfg(feature = "metrics")]
 pub use crate::metric::classification::ClassReduction;
 // Expose `ItemLazy` so it can be implemented for custom types
 pub use processor::ItemLazy;

crates/burn/Cargo.toml

@@ -24,7 +24,7 @@ train = ["burn-train", "autodiff", "dataset"]
 tui = ["burn-train?/tui"]
 
 ##  Includes system info metrics (CPU/GPU usage, etc)
-metrics = ["burn-train?/metrics"]
+metrics = ["burn-train?/sys-metrics"]
 
 # Datasets
 dataset = ["burn-core/dataset"]