TrainingLoop: refactor progress printer and add CSVLogger

Examples/LeNet-MNIST/main.swift

@@ -43,12 +43,14 @@ var classifier = Sequential {
 
 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(liveStatistics: false).log])
+trainingLoop.statisticsRecorder.liveStatistics = false
+trainingLoop.progressPrinter.liveStatistics = false
 
 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,79 @@
+import Foundation
+import ModelSupport
+
+/// A callback-based handler for logging the statistics to CSV file.
+public class CSVLogger {
+  public var path: String
+  public var liveStatistics: Bool
+
+  let foundationFS: FoundationFileSystem
+  let foundationFile: FoundationFile
+
+  /// Create an instance that log statistics during the training loop.
+  ///
+  /// - Parameters:
+  ///   - liveStatistics: whether or not log the statistics lively on each batch.
+  public init(withPath path: String = "run/log.csv", liveStatistics: Bool = true) {
+    self.path = path
+    self.liveStatistics = liveStatistics
+    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 {
+        break
+      }
+
+      if !liveStatistics && (batchIndex + 1 != batchCount) {
+        break
+      }
+
+      guard let stats = loop.lastStatsLog else {
+        break
+      }
+
+      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:
+      break
+    }
+  }
+
+  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,100 @@
+// 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 {
+  public var liveStatistics: Bool
+
+  /// Create an instance that prints progress during the training loop.
+  /// The progress contains a dynamic progress bar followed by statistics of metrics.
+  ///
+  /// - Parameters:
+  ///   - liveStatistics: whether or not show the statistics lively on each batch.
+  public init(liveStatistics: Bool = true) {
+    self.liveStatistics = liveStatistics
+  }
+
+  /// 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 liveStatistics || (batchCount == (batchIndex + 1)) {
+        stats = formatStats(loop.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 = ""
+    if let stats = stats {
+      for stat in stats {
+        result += " - \(stat.0): \(String(format: "%.4f", stat.1))"
+      }
+    }
+    return result
+  }
+}

TrainingLoop/Callbacks/StatisticsRecorder.swift

@@ -0,0 +1,87 @@
+// 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 TensorFlow
+
+/// A callback-based handler for recording statistics.
+///
+/// Data produced by this handler can be used by ProgressPrinter, CVSLogger, etc.
+public class StatisticsRecorder {
+  public var liveStatistics: Bool
+
+  var metricMeasurers: [MetricsMeasurer]
+
+  /// Create an instance that records 'metrics' during the training loop.
+  /// 
+  /// Recording happens every batch by default or
+  /// only when last batch ends when 'liveStatistics' is set to false.
+  ///
+  /// - Parameters:
+  ///   - liveStatistics: whether or not record lively on each batch.
+  ///     This has an impact on performance, due to materialization of tensors, updating values
+  ///     on every batch can reduce training speed by up to 30%.
+  ///   - metrics: an array of TrainingMetrics to record.
+  public init(liveStatistics: Bool = true, metrics: [TrainingMetrics]) {
+    self.liveStatistics = liveStatistics
+    metricMeasurers = metrics.map { $0.measurer }
+  }
+
+  /// The callback used to hook into the TrainingLoop for recording statistics.
+  ///
+  /// It will record the statistics into lastStatsLog in the loop where other 
+  /// callbacks can consume from.
+  ///
+  /// - Parameters:
+  ///   - loop: The TrainingLoop where an event has occurred. 
+  ///   - event: The training or validation event that this callback is responding to.
+  public func record<L: TrainingLoopProtocol>(_ loop: inout L, event: TrainingLoopEvent) throws {
+    switch event {
+    case .trainingStart, .validationStart:
+      resetMetricMeasurers()
+    case .batchEnd:
+      if let loss = loop.lastStepLoss, let output = loop.lastStepOutput,
+        let target = loop.lastStepTarget
+      {
+        accumulateMetrics(loss: loss, predictions: output, labels: target)
+      }
+
+      if let batchIndex = loop.batchIndex, let batchCount = loop.batchCount {
+        if liveStatistics || (batchCount == (batchIndex + 1)) {
+          loop.lastStatsLog = computeMetrics()
+        }
+      }
+    default:
+      return
+    }
+  }
+
+  func resetMetricMeasurers() {
+    for index in metricMeasurers.indices {
+      metricMeasurers[index].reset()
+    }
+  }
+
+  func accumulateMetrics<Output, Target>(loss: Tensor<Float>, predictions: Output, labels: Target) {
+    for index in metricMeasurers.indices {
+      metricMeasurers[index].accumulate(loss: loss, predictions: predictions, labels: labels)
+    }
+  }
+
+  func computeMetrics() -> [(String, Float)] {
+    var result: [(String, Float)] = []
+    for measurer in metricMeasurers {
+      result.append((measurer.name, measurer.measure()))
+    }
+    return result
+  }
+}