TrainingLoop: refactor progress printer and add CSVLogger

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: [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,11 @@ public struct FoundationFile: File {
     // TODO: Incorporate file offset.
     try value.write(to: location)
   }
+
+  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,70 @@
+import Foundation
+import ModelSupport
+
+/// A callback-based handler for logging the statistics to CSV file.
+public class CSVLogger {
+  public var path: String
+
+  let foundationFS: FoundationFileSystem
+  let foundationFile: FoundationFile
+
+  /// Create an instance that log statistics during the training loop.
+  public init(withPath path: String = "run/log.csv") {
+    self.path = path
+    self.foundationFS = FoundationFileSystem()
+    self.foundationFile = FoundationFile(path: path)
+  }
+
+  /// The callback used to hook into the TrainingLoop for logging statistics.
+  ///
+  /// - Parameters:
+  ///   - loop: The TrainingLoop where an event has occurred. 
+  ///   - event: The training or validation event that this callback is responding to.
+  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
+      else {
+        return
+      }
+
+      guard let stats = loop.lastStatsLog else {
+        return
+      }
+
+      if !FileManager.default.fileExists(atPath: path) {
+        try foundationFS.createDirectoryIfMissing(at: String(path[..<path.lastIndex(of: "/")!]))
+        try writeHeader(stats: stats)
+      }
+      try writeDataRow(
+        epoch: "\(epochIndex + 1)/\(epochCount)",
+        batch: "\(batchIndex + 1)/\(batchCount)",
+        stats: stats)
+    default:
+      return
+    }
+  }
+
+  func writeHeader(stats: [(String, Float)]) throws {
+    let head: String = (["epoch", "batch"] + stats.map { $0.0 }).joined(separator: ", ")
+    do {
+      try head.write(toFile: path, atomically: true, encoding: .utf8)
+    } catch {
+      print("Unexpected error in writing header line: \(error).")
+      throw error
+    }
+  }
+
+  func writeDataRow(epoch: String, batch: String, stats: [(String, Float)]) throws {
+    let dataRow: Data = (
+      "\n" + ([epoch, batch] + stats.map { String($0.1) }).joined(separator: ", ")
+    ).data(using: .utf8)!
+    do {
+      try foundationFile.append(dataRow)
+    } catch {
+      print("Unexpected error in writing data row: \(error).")
+      throw error
+    }
+  }
+}

TrainingLoop/Callbacks/ProgressPrinter.swift

@@ -0,0 +1,93 @@
+// 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 callback-based handler for printing the training or validation progress. 
+public class ProgressPrinter {
+
+  /// Create an instance that prints progress during the training loop.
+  /// The progress contains a dynamic progress bar followed by statistics of metrics.
+  public init() {
+  }
+
+  /// The callback used to hook into the TrainingLoop for printing progress.
+  /// 
+  /// 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
+  ///
+  /// - Parameters:
+  ///   - loop: The TrainingLoop where an event has occurred. 
+  ///   - event: The training or validation event that this callback is responding to.
+  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 {
+        return
+      }
+
+      Swift.print("Epoch \(epochIndex + 1)/\(epochCount)")
+    case .batchEnd:
+      guard let batchIndex = loop.batchIndex, let batchCount = loop.batchCount else {
+        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
+  }
+}