refactor: multi-thread logger -> async logger

djdisodo · Aug 27, 2022 · 9ed0b0b · 9ed0b0b
1 parent 05eb2f9
commit 9ed0b0b
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/examples/mnist.rs b/examples/mnist.rs
@@ -8,7 +8,7 @@ use burn::tensor::af::relu;
 use burn::tensor::back::{ad, Backend};
 use burn::tensor::losses::cross_entropy_with_logits;
 use burn::tensor::{Data, ElementConversion, Shape, Tensor};
-use burn::train::logger::{CLILogger, MultiThreadLogger};
+use burn::train::logger::{AsyncLogger, CLILogger};
 use burn::train::metric::{AccuracyMetric, CUDAMetric, LossMetric, Metric};
 use burn::train::{ClassificationLearner, ClassificationOutput, SupervisedTrainer};
 use std::sync::Arc;
@@ -180,11 +180,11 @@ fn run<B: ad::Backend>(device: B::Device) {
 
     let learner = ClassificationLearner::new(model);
 
-    let logger_train = Box::new(MultiThreadLogger::new(Box::new(CLILogger::new(
+    let logger_train = Box::new(AsyncLogger::new(Box::new(CLILogger::new(
         metrics(),
         "Train".to_string(),
     ))));
-    let logger_valid = Box::new(MultiThreadLogger::new(Box::new(CLILogger::new(
+    let logger_valid = Box::new(AsyncLogger::new(Box::new(CLILogger::new(
         metrics(),
         "Valid".to_string(),
     ))));

diff --git a/src/train/logger/multi_thread.rs → src/train/logger/async_logger.rs b/src/train/logger/multi_thread.rs → src/train/logger/async_logger.rs
@@ -6,7 +6,7 @@ enum Message<T> {
     Clear,
 }
 
-pub struct MultiThreadLogger<T> {
+pub struct AsyncLogger<T> {
     sender: mpsc::Sender<Message<T>>,
 }
 
@@ -33,7 +33,7 @@ impl<T> LoggerThread<T> {
     }
 }
 
-impl<T: Send + Sync + 'static> MultiThreadLogger<T> {
+impl<T: Send + Sync + 'static> AsyncLogger<T> {
     pub fn new(logger: Box<dyn Logger<T>>) -> Self {
         let (sender, receiver) = mpsc::channel();
         let thread = LoggerThread::new(Mutex::new(logger), receiver);
@@ -44,7 +44,7 @@ impl<T: Send + Sync + 'static> MultiThreadLogger<T> {
     }
 }
 
-impl<T: Send> Logger<T> for MultiThreadLogger<T> {
+impl<T: Send> Logger<T> for AsyncLogger<T> {
     fn log(&mut self, item: T) {
         self.sender.send(Message::Log(item)).unwrap();
     }

diff --git a/src/train/logger/mod.rs b/src/train/logger/mod.rs
@@ -1,7 +1,7 @@
+mod async_logger;
 mod cli;
 mod logger;
-mod multi_thread;
 
+pub use async_logger::*;
 pub use cli::*;
 pub use logger::*;
-pub use multi_thread::*;