tensorboard_rs/
summary_writer.rs

1#![allow(clippy::too_many_arguments)]
2use std::path::{PathBuf, Path};
3use std::time::SystemTime;
4use std::collections::HashMap;
5use protobuf::Message;
6use protobuf::RepeatedField;
7use crate::event_file_writer::EventFileWriter;
8use tensorboard_proto::event::{Event, TaggedRunMetadata};
9use tensorboard_proto::summary::{Summary};
10use tensorboard_proto::graph::{GraphDef, };
11use tensorboard_proto::node_def::{NodeDef, };
12use tensorboard_proto::versions::{VersionDef, };
13//use tensorboard_proto::attr_value::{AttrValue, };
14//use tensorboard_proto::tensor_shape::{TensorShapeProto, };
15use tensorboard_proto::step_stats::{RunMetadata, };
16use crate::summary::{scalar, image, histogram_raw};
17
18
19pub struct FileWriter {
20    writer: EventFileWriter,
21}
22impl FileWriter {
23    pub fn new<P: AsRef<Path>>(logdir: P) -> FileWriter {
24        FileWriter {
25            writer: EventFileWriter::new(logdir),
26        }
27    }
28    pub fn get_logdir(&self) -> PathBuf {
29        self.writer.get_logdir()
30    }
31    pub fn add_event(&mut self, event: &Event, step: usize) {
32        let mut event = event.clone();
33        
34        let mut time_full = 0.0;
35        if let Ok(n) = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) {
36            time_full = n.as_secs_f64();
37        }
38        event.set_wall_time(time_full);
39        
40        event.set_step(step as i64);
41        
42        self.writer.add_event(&event)
43    }
44    pub fn add_summary(&mut self, summary: Summary, step: usize) {
45        let mut evn = Event::new();
46        evn.set_summary(summary);
47        self.add_event(&evn, step)
48    }
49    pub fn add_graph(&mut self, graph: GraphDef, meta: RunMetadata) {
50        let mut graph_vec: Vec<u8> = Vec::new();
51        graph.write_to_vec(&mut graph_vec).expect("");
52        let mut graph_evn = Event::new();
53        graph_evn.set_graph_def(graph_vec);
54        self.writer.add_event(&graph_evn);
55
56        let mut meta_vec: Vec<u8> = Vec::new();
57        meta.write_to_vec(&mut meta_vec).expect("");
58        let mut tagged_meta = TaggedRunMetadata::new();
59        tagged_meta.set_tag("profiler".to_string());
60        tagged_meta.set_run_metadata(meta_vec);
61        let mut meta_evn = Event::new();
62        meta_evn.set_tagged_run_metadata(tagged_meta);
63        self.writer.add_event(&meta_evn);
64    }
65    pub fn flush(&mut self) {
66        self.writer.flush()
67    }
68}
69
70pub struct SummaryWriter {
71    writer: FileWriter,
72    all_writers: HashMap<PathBuf, FileWriter>,
73}
74impl SummaryWriter {
75    pub fn new<P: AsRef<Path>>(logdir: P) -> SummaryWriter {
76        SummaryWriter {
77            writer: FileWriter::new(logdir),
78            all_writers: HashMap::new(),
79        }
80    }
81    pub fn add_hparams(&mut self) {unimplemented!();}
82    pub fn add_scalar(&mut self, tag: &str, scalar_value: f32, step: usize) {
83        self.writer.add_summary(scalar(tag, scalar_value), step);
84    }
85    pub fn add_scalars(&mut self, main_tag: &str, tag_scalar: &HashMap<String, f32>, step: usize) {
86        let base_logdir = self.writer.get_logdir();
87        for (tag, scalar_value) in tag_scalar.iter() {
88            let fw_tag = base_logdir.join(main_tag).join(tag);
89            if ! self.all_writers.contains_key(&fw_tag) {
90                let new_writer = FileWriter::new(fw_tag.clone());
91                self.all_writers.insert(fw_tag.clone(), new_writer);
92            }
93            let fw = self.all_writers.get_mut(&fw_tag).expect("");
94            fw.add_summary(scalar(main_tag, *scalar_value), step);
95        }
96    }
97
98    pub fn export_scalars_to_json(&self) {unimplemented!();}
99    pub fn add_histogram(&mut self) {unimplemented!();}
100    pub fn add_histogram_raw(&mut self,
101                             tag: &str,
102                             min: f64, max: f64,
103                             num: f64,
104                             sum: f64, sum_squares: f64,
105                             bucket_limits: &[f64], bucket_counts: &[f64],
106                             step: usize
107    ) {
108        if bucket_limits.len() != bucket_counts.len() {
109            panic!("bucket_limits.len() != bucket_counts.len()");
110        }
111
112        self.writer.add_summary(histogram_raw(tag, min, max, num, sum, sum_squares, bucket_limits, bucket_counts), step);
113    }
114    pub fn add_image(&mut self, tag: &str, data: &[u8], dim: &[usize], step: usize) {
115        self.writer.add_summary(image(tag, data, dim), step);
116    }
117    pub fn add_images(&mut self) {unimplemented!();}
118    pub fn add_image_with_boxes(&mut self) {unimplemented!();}
119    pub fn add_figure(&mut self) {unimplemented!();}
120    pub fn add_video(&mut self) {unimplemented!();}
121    pub fn add_audio(&mut self) {unimplemented!();}
122    pub fn add_text(&mut self) {unimplemented!();}
123    pub fn add_onnx_graph(&mut self) {unimplemented!();}
124    pub fn add_openvino_graph(&mut self) {unimplemented!();}
125    pub fn add_graph(&mut self, node_list: &[NodeDef]) {
126        let mut graph = GraphDef::new();
127        
128        let nodes = RepeatedField::from(node_list.to_vec());
129        graph.set_node(nodes);
130        
131        let mut version = VersionDef::new();
132        version.set_producer(22);
133        graph.set_versions(version);
134
135        let stats = RunMetadata::new();
136        
137        self.writer.add_graph(graph, stats);
138    }
139    pub fn add_embedding(&mut self) {unimplemented!();}
140    pub fn add_pr_curve(&mut self) {unimplemented!();}
141    pub fn add_pr_curve_raw(&mut self) {unimplemented!();}
142    pub fn add_custom_scalars_multilinechart(&mut self) {unimplemented!();}
143    pub fn add_custom_scalars_marginchart(&mut self) {unimplemented!();}
144    pub fn add_custom_scalars(&mut self) {unimplemented!();}
145    pub fn add_mesh(&mut self) {unimplemented!();}
146
147    pub fn flush(&mut self) {
148        self.writer.flush();
149    }
150}
tensorboard_rs/summary_writer.rs

tensorboard_rs/
summary_writer.rs
