burn_train/metric/
memory_use.rs

1/// RAM use metric
2use super::{MetricMetadata, Numeric};
3use crate::metric::{Metric, MetricEntry};
4use std::time::{Duration, Instant};
5use sysinfo::System;
6
7/// Memory information
8pub struct CpuMemory {
9    last_refresh: Instant,
10    refresh_frequency: Duration,
11    sys: System,
12    ram_bytes_total: u64,
13    ram_bytes_used: u64,
14}
15
16impl CpuMemory {
17    /// Creates a new memory metric
18    pub fn new() -> Self {
19        let mut metric = Self {
20            last_refresh: Instant::now(),
21            refresh_frequency: Duration::from_millis(200),
22            sys: System::new(),
23            ram_bytes_total: 0,
24            ram_bytes_used: 0,
25        };
26        metric.refresh();
27        metric
28    }
29
30    fn refresh(&mut self) {
31        self.sys.refresh_memory();
32        self.last_refresh = Instant::now();
33
34        // bytes of RAM available
35        self.ram_bytes_total = self.sys.total_memory();
36
37        // bytes of RAM in use
38        self.ram_bytes_used = self.sys.used_memory();
39    }
40}
41
42impl Default for CpuMemory {
43    fn default() -> Self {
44        CpuMemory::new()
45    }
46}
47
48impl Metric for CpuMemory {
49    type Input = ();
50
51    fn update(&mut self, _item: &Self::Input, _metadata: &MetricMetadata) -> MetricEntry {
52        if self.last_refresh.elapsed() >= self.refresh_frequency {
53            self.refresh();
54        }
55
56        let raw = bytes2gb(self.ram_bytes_used);
57        let formatted = format!(
58            "RAM Used: {:.2} / {:.2} Gb",
59            raw,
60            bytes2gb(self.ram_bytes_total),
61        );
62
63        MetricEntry::new(self.name(), formatted, raw.to_string())
64    }
65
66    fn clear(&mut self) {}
67
68    fn name(&self) -> String {
69        "CPU Memory".to_string()
70    }
71}
72
73impl Numeric for CpuMemory {
74    fn value(&self) -> f64 {
75        bytes2gb(self.ram_bytes_used)
76    }
77}
78
79fn bytes2gb(bytes: u64) -> f64 {
80    bytes as f64 / 1e9
81}
burn_train/metric/memory_use.rs

burn_train/metric/
memory_use.rs
