DistriOptimizerV2 logger (intel-analytics#3135)

* DistriOptimizerV2 logger * update * fix style check * validate epoch num
dding3 · Jun 30, 2021 · 284ded4 · 284ded4
1 parent 39cb177
commit 284ded4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 16 deletions.
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizerV2.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizerV2.scala
@@ -56,16 +56,7 @@ object DistriOptimizerV2 extends AbstractOptimizer {
     parameterProcessers: Array[ParameterProcessor] = null) extends DistriOptimizer.Cache[T]
 
   import Optimizer._
-
-  private[DistriOptimizerV2] var _logger: Option[OptimizerLogger] = None
-
-  def logger: OptimizerLogger = {
-    if (_logger.isEmpty) {
-      _logger = Some(new DistriLogger)
-    }
-
-    _logger.get
-  }
+  val logger: Logger = Logger.getLogger(getClass)
 
   private[optim] def optimize[T: ClassTag](
     cacheOfMaster: MasterCache[T],
@@ -107,7 +98,7 @@ object DistriOptimizerV2 extends AbstractOptimizer {
       }
 
       val _header = header(
-        trainingTrace.epochs,
+        trainingTrace.epochs-1,
         trainingTrace.recordsOfEpoch,
         context.numSamples,
         trainingTrace.iterations,
@@ -857,10 +848,6 @@ class DistriOptimizerV2[T: ClassTag](
     CachedModels.deleteKey(modelBroadcast.uuid)
   }
 
-  def setLogger(logger: OptimizerLogger): Unit = {
-    DistriOptimizerV2._logger = Some(logger)
-  }
-
   private def validArgs(): Boolean = {
     val checkSingleton = this.checkSingleton
     val nodeNumber = Engine.nodeNumber()

diff --git a/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LoggerFilterSpec.scala b/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LoggerFilterSpec.scala
@@ -121,7 +121,7 @@ class LoggerFilterSpec extends FlatSpec with BeforeAndAfter with Matchers {
     }
 
     {
-      val pattern = s".*INFO.*DistriLogger.* - " + "" +
+      val pattern = s".*INFO.*DistriOptimizerV2.* - " + "" +
         s"\\[Epoch 1 100/100\\]\\[Iteration 2\\]\\[Wall Clock .*\\] " +
         s"Epoch finished. Wall clock time is .*ms"