burn 0.3.0

BURN: Burn Unstoppable Rusty Neurons
Documentation
use super::RunningMetricResult;
use crate::train::metric::{Metric, MetricState};
use nvml_wrapper::Nvml;

pub struct CUDAMetric {
    nvml: Nvml,
}

impl CUDAMetric {
    pub fn new() -> Self {
        Self {
            nvml: Nvml::init().unwrap(),
        }
    }
}

impl Default for CUDAMetric {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Metric<T> for CUDAMetric {
    fn update(&mut self, _item: &T) -> Box<dyn MetricState> {
        let name = String::from("Cuda");

        let mut formatted = String::new();
        let mut raw_running = String::new();

        for index in 0..self.nvml.device_count().unwrap() {
            let device = self.nvml.device_by_index(index).unwrap();
            let memory_info = device.memory_info().unwrap();
            let used_gb = memory_info.used as f64 * 1e-9;
            let total_gb = memory_info.total as f64 * 1e-9;

            let memory_info_formatted = format!("{:.2}/{:.2} Gb", used_gb, total_gb);
            let memory_info_raw = format!("{}/{}", used_gb, total_gb);

            formatted = format!(
                "{} GPU #{} - Memory {}",
                formatted, index, memory_info_formatted
            );
            raw_running = format!("{} ", memory_info_raw);

            let utilization_rates = device.utilization_rates().unwrap();
            let utilization_rate_formatted = format!("{}%", utilization_rates.gpu);
            formatted = format!("{} - Usage {}", formatted, utilization_rate_formatted);
        }

        Box::new(RunningMetricResult {
            name,
            formatted,
            raw_running,
            raw_current: String::new(),
        })
    }

    fn clear(&mut self) {}
}