TrainingLoop: refactor progress printer and add CSVLogger (#668)

Author: xihui-wu

Examples/LeNet-MNIST/main.swift

@@ -31,24 +31,33 @@ let dataset = MNIST(batchSize: batchSize, on: device)
 
 // The LeNet-5 model, equivalent to `LeNet` in `ImageClassificationModels`.
 var classifier = Sequential {
-    Conv2D<Float>(filterShape: (5, 5, 1, 6), padding: .same, activation: relu)
-    AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
-    Conv2D<Float>(filterShape: (5, 5, 6, 16), activation: relu)
-    AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
-    Flatten<Float>()
-    Dense<Float>(inputSize: 400, outputSize: 120, activation: relu)
-    Dense<Float>(inputSize: 120, outputSize: 84, activation: relu)
-    Dense<Float>(inputSize: 84, outputSize: 10)
+  Conv2D<Float>(filterShape: (5, 5, 1, 6), padding: .same, activation: relu)
+  AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
+  Conv2D<Float>(filterShape: (5, 5, 6, 16), activation: relu)
+  AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
+  Flatten<Float>()
+  Dense<Float>(inputSize: 400, outputSize: 120, activation: relu)
+  Dense<Float>(inputSize: 120, outputSize: 84, activation: relu)
+  Dense<Float>(inputSize: 84, outputSize: 10)
 }
 
 var optimizer = SGD(for: classifier, learningRate: 0.1)
 
-let trainingProgress = TrainingProgress()
 var trainingLoop = TrainingLoop(
   training: dataset.training,
   validation: dataset.validation,
   optimizer: optimizer,
   lossFunction: softmaxCrossEntropy,
-  callbacks: [trainingProgress.update])
+  metrics: [.accuracy],
+  callbacks: [try! CSVLogger().log])
+
+// Compute statistics only when last batch ends.
+trainingLoop.statisticsRecorder!.shouldCompute = {
+  (
+    _ batchIndex: Int, _ batchCount: Int, _ epochIndex: Int, _ epochCount: Int,
+    _ event: TrainingLoopEvent
+  ) -> Bool in
+  return event == .batchEnd && batchIndex + 1 == batchCount
+}
 
 try! trainingLoop.fit(&classifier, epochs: epochCount, on: device)

Examples/MobileNetV1-Imagenette/main.swift

@@ -29,12 +29,11 @@ let dataset = Imagenette(batchSize: 64, inputSize: .resized320, outputSize: 224,
 var model = MobileNetV1(classCount: 10)
 let optimizer = SGD(for: model, learningRate: 0.02, momentum: 0.9)
 
-let trainingProgress = TrainingProgress()
 var trainingLoop = TrainingLoop(
   training: dataset.training,
   validation: dataset.validation,
   optimizer: optimizer,
   lossFunction: softmaxCrossEntropy,
-  callbacks: [trainingProgress.update])
+  metrics: [.accuracy])
 
 try! trainingLoop.fit(&model, epochs: 10, on: device)

Examples/MobileNetV2-Imagenette/main.swift

@@ -29,12 +29,11 @@ let dataset = Imagenette(batchSize: 64, inputSize: .resized320, outputSize: 224,
 var model = MobileNetV2(classCount: 10)
 let optimizer = SGD(for: model, learningRate: 0.002, momentum: 0.9)
 
-let trainingProgress = TrainingProgress()
 var trainingLoop = TrainingLoop(
   training: dataset.training,
   validation: dataset.validation,
   optimizer: optimizer,
   lossFunction: softmaxCrossEntropy,
-  callbacks: [trainingProgress.update])
+  metrics: [.accuracy])
 
 try! trainingLoop.fit(&model, epochs: 10, on: device)

Examples/ResNet-CIFAR10/main.swift

@@ -29,12 +29,11 @@ let dataset = CIFAR10(batchSize: 10, on: device)
 var model = ResNet(classCount: 10, depth: .resNet56, downsamplingInFirstStage: false)
 var optimizer = SGD(for: model, learningRate: 0.001)
 
-let trainingProgress = TrainingProgress()
 var trainingLoop = TrainingLoop(
   training: dataset.training,
   validation: dataset.validation,
   optimizer: optimizer,
   lossFunction: softmaxCrossEntropy,
-  callbacks: [trainingProgress.update])
+  metrics: [.accuracy])
 
 try! trainingLoop.fit(&model, epochs: 10, on: device)

Examples/VGG-Imagewoof/main.swift

@@ -39,12 +39,12 @@ public func scheduleLearningRate<L: TrainingLoopProtocol>(
   }
 }
 
-let trainingProgress = TrainingProgress()
 var trainingLoop = TrainingLoop(
   training: dataset.training,
   validation: dataset.validation,
   optimizer: optimizer,
   lossFunction: softmaxCrossEntropy,
-  callbacks: [trainingProgress.update, scheduleLearningRate])
+  metrics: [.accuracy],
+  callbacks: [scheduleLearningRate])
 
 try! trainingLoop.fit(&model, epochs: 90, on: device)

Support/FileSystem.swift

@@ -39,4 +39,5 @@ public protocol File {
   func read(position: Int, count: Int) throws -> Data
   func write(_ value: Data) throws
   func write(_ value: Data, position: Int) throws
+  func append(_ value: Data) throws
 }

Support/FoundationFileSystem.swift

@@ -58,4 +58,14 @@ public struct FoundationFile: File {
     // TODO: Incorporate file offset.
     try value.write(to: location)
   }
+
+  /// Append data to the file.
+  ///
+  /// Parameter value: data to be appended at the end.
+  public func append(_ value: Data) throws {
+    let fileHandler = try FileHandle(forUpdating: location)
+    try fileHandler.seekToEnd()
+    try fileHandler.write(contentsOf: value)
+    try fileHandler.close()
+  }
 }

TrainingLoop/CMakeLists.txt

@@ -1,8 +1,10 @@
 add_library(TrainingLoop
   LossFunctions.swift
+  Metrics.swift
   TrainingLoop.swift
-  TrainingProgress.swift
-  TrainingStatistics.swift)
+  Callbacks/StatisticsRecorder.swift
+  Callbacks/ProgressPrinter.swift
+  Callbacks/CSVLogger.swift)
 target_link_libraries(TrainingLoop PUBLIC
   ModelSupport)
 set_target_properties(TrainingLoop PROPERTIES

TrainingLoop/Callbacks/CSVLogger.swift

@@ -0,0 +1,74 @@
+import Foundation
+import ModelSupport
+
+public enum CSVLoggerError: Error {
+  case InvalidPath
+}
+
+/// A handler for logging training and validation statistics to a CSV file.
+public class CSVLogger {
+  /// The path of the file that statistics are logged to.
+  public var path: String
+
+  // True iff the header of the CSV file has been written.
+  fileprivate var headerWritten: Bool
+
+  /// Creates an instance that logs to a file with the given path.
+  ///
+  /// Throws: File system errors.
+  public init(path: String = "run/log.csv") throws {
+    self.path = path
+
+    // Validate the path.
+    let url = URL(fileURLWithPath: path)
+    if url.pathExtension != "csv" {
+      throw CSVLoggerError.InvalidPath
+    }
+    // Create the containing directory if it is missing.
+    try FoundationFileSystem().createDirectoryIfMissing(at: url.deletingLastPathComponent().path)
+    // Initialize the file with empty string.
+    try FoundationFile(path: path).write(Data())
+
+    self.headerWritten = false
+  }
+
+  /// Logs the statistics for the 'loop' when 'batchEnd' event happens; 
+  /// ignoring other events.
+  ///
+  /// Throws: File system errors.
+  public func log<L: TrainingLoopProtocol>(_ loop: inout L, event: TrainingLoopEvent) throws {
+    switch event {
+    case .batchEnd:
+      guard let epochIndex = loop.epochIndex, let epochCount = loop.epochCount,
+        let batchIndex = loop.batchIndex, let batchCount = loop.batchCount,
+        let stats = loop.lastStatsLog
+      else {
+        // No-Op if trainingLoop doesn't set the required values for stats logging.
+        return
+      }
+
+      if !headerWritten {
+        try writeHeader(stats: stats)
+        headerWritten = true
+      }
+
+      try writeDataRow(
+        epoch: "\(epochIndex + 1)/\(epochCount)",
+        batch: "\(batchIndex + 1)/\(batchCount)",
+        stats: stats)
+    default:
+      return
+    }
+  }
+
+  func writeHeader(stats: [(name: String, value: Float)]) throws {
+    let header = (["epoch", "batch"] + stats.lazy.map { $0.name }).joined(separator: ", ") + "\n"
+    try FoundationFile(path: path).append(header.data(using: .utf8)!)
+  }
+
+  func writeDataRow(epoch: String, batch: String, stats: [(name: String, value: Float)]) throws {
+    let dataRow = ([epoch, batch] + stats.lazy.map { String($0.value) }).joined(separator: ", ")
+      + "\n"
+    try FoundationFile(path: path).append(dataRow.data(using: .utf8)!)
+  }
+}

TrainingLoop/Callbacks/ProgressPrinter.swift

@@ -0,0 +1,85 @@
+// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+let progressBarLength = 30
+
+/// A handler for printing the training and validation progress. 
+public class ProgressPrinter {
+  /// Print training or validation progress in response of the 'event'.
+  /// 
+  /// An example of the progress would be:
+  /// Epoch 1/12
+  /// 468/468 [==============================] - loss: 0.4819 - accuracy: 0.8513
+  /// 79/79 [==============================] - loss: 0.1520 - accuracy: 0.9521
+  public func print<L: TrainingLoopProtocol>(_ loop: inout L, event: TrainingLoopEvent) throws {
+    switch event {
+    case .epochStart:
+      guard let epochIndex = loop.epochIndex, let epochCount = loop.epochCount else {
+        // No-Op if trainingLoop doesn't set the required values for progress printing.
+        return
+      }
+
+      Swift.print("Epoch \(epochIndex + 1)/\(epochCount)")
+    case .batchEnd:
+      guard let batchIndex = loop.batchIndex, let batchCount = loop.batchCount else {
+        // No-Op if trainingLoop doesn't set the required values for progress printing.
+        return
+      }
+
+      let progressBar = formatProgressBar(
+        progress: Float(batchIndex + 1) / Float(batchCount), length: progressBarLength)
+      var stats: String = ""
+      if let lastStatsLog = loop.lastStatsLog {
+        stats = formatStats(lastStatsLog)
+      }
+
+      Swift.print(
+        "\r\(batchIndex + 1)/\(batchCount) \(progressBar)\(stats)",
+        terminator: ""
+      )
+      fflush(stdout)
+    case .epochEnd:
+      Swift.print("")
+    case .validationStart:
+      Swift.print("")
+    default:
+      return
+    }
+  }
+
+  func formatProgressBar(progress: Float, length: Int) -> String {
+    let progressSteps = Int(round(Float(length) * progress))
+    let leading = String(repeating: "=", count: progressSteps)
+    let separator: String
+    let trailing: String
+    if progressSteps < progressBarLength {
+      separator = ">"
+      trailing = String(repeating: ".", count: progressBarLength - progressSteps - 1)
+    } else {
+      separator = ""
+      trailing = ""
+    }
+    return "[\(leading)\(separator)\(trailing)]"
+  }
+
+  func formatStats(_ stats: [(String, Float)]) -> String {
+    var result = ""
+    for stat in stats {
+      result += " - \(stat.0): \(String(format: "%.4f", stat.1))"
+    }
+    return result
+  }
+}