burn_train/metric/
memory_use.rs

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

burn_train/metric/
memory_use.rs
