Restructuring checkpoint files, adding an initial filesystem abstraction (#618)

* Initial move of the checkpoint reader / writer to their own module.

* Updating BERT for new checkpoint organization.

* Update MiniGo for checkpoint migration.

* Initial experiments with a filesystem abstraction.

* Using a default filesystem for the index reader.

* Minor formatting.

* Remove unnecessary default.

* Remove CMake references to old Checkpoint files.

* Fixing more CMake imports.

* Actually add the Checkpoints CMakelists.

* Enable testing and proper dependency for CMake.

* Adding the last few CMake files.

Author: BradLarson

CMakeLists.txt

@@ -124,6 +124,7 @@ add_dependencies(ArgumentParser swift-argument-parser-install)
 add_subdirectory(Autoencoder)
 add_subdirectory(Support)
 add_subdirectory(Batcher)
+add_subdirectory(Checkpoints)
 add_subdirectory(Datasets)
 add_subdirectory(Models)
 add_subdirectory(MiniGo)

Checkpoints/CMakeLists.txt

@@ -0,0 +1,27 @@
+add_library(Checkpoints
+  CheckpointIndexReader.swift
+  CheckpointIndexWriter.swift
+  CheckpointReader.swift
+  CheckpointWriter.swift
+  Protobufs/tensor_bundle.pb.swift
+  Protobufs/tensor_shape.pb.swift
+  Protobufs/tensor_slice.pb.swift
+  Protobufs/types.pb.swift
+  Protobufs/versions.pb.swift
+  SnappyDecompression.swift)
+set_target_properties(Checkpoints PROPERTIES
+  INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})
+target_compile_options(Checkpoints PRIVATE
+  $<$<BOOL:${BUILD_TESTING}>:-enable-testing>)
+target_link_libraries(Checkpoints PUBLIC
+  ModelSupport)
+
+install(TARGETS Checkpoints
+  ARCHIVE DESTINATION lib/swift/$<LOWER_CASE:${CMAKE_SYSTEM_NAME}>
+  LIBRARY DESTINATION lib/swift/$<LOWER_CASE:${CMAKE_SYSTEM_NAME}>
+  RUNTIME DESTINATION bin)
+get_swift_host_arch(swift_arch)
+install(FILES
+  $<TARGET_PROPERTY:Checkpoints,Swift_MODULE_DIRECTORY>/Checkpoints.swiftdoc
+  $<TARGET_PROPERTY:Checkpoints,Swift_MODULE_DIRECTORY>/Checkpoints.swiftmodule
+  DESTINATION lib/swift$<$<NOT:$<BOOL:${BUILD_SHARED_LIBS}>>:_static>/$<LOWER_CASE:${CMAKE_SYSTEM_NAME}>/${swift_arch})

Support/Checkpoints/CheckpointIndexReader.swift Checkpoints/CheckpointIndexReader.swift

@@ -30,6 +30,7 @@
 // locations determined by this metadata.
 
 import Foundation
+import ModelSupport
 
 // The block footer size is constant, and is obtained from the following:
 // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/lib/io/format.h
@@ -43,8 +44,9 @@ class CheckpointIndexReader {
 
     var atEndOfFile: Bool { return index >= (binaryData.count - footerSize - 1) }
 
-    init(file: URL) throws {
-        let fileData = try Data(contentsOf: file)
+    init(file: URL, fileSystem: FileSystem = FoundationFileSystem()) throws {
+        let indexFile = fileSystem.open(file.path)
+        let fileData = try indexFile.read()
         if fileData[0] == 0 {
             binaryData = fileData
         } else {

Support/Checkpoints/CheckpointIndexWriter.swift Checkpoints/CheckpointIndexWriter.swift

@@ -20,6 +20,7 @@
 // shards to obtain their raw bytes.
 
 import Foundation
+import ModelSupport
 import TensorFlow
 
 /// A Swift-native TensorFlow v2 checkpoint reader that can download all checkpoint files from 
@@ -29,6 +30,7 @@ open class CheckpointReader {
     let header: Tensorflow_BundleHeaderProto
     let metadata: [String: Tensorflow_BundleEntryProto]
     var shardCache: [URL: Data] = [:]
+    let fileSystem: FileSystem
 
     /// The local checkpoint location.
     public let localCheckpointLocation: URL
@@ -51,7 +53,11 @@ open class CheckpointReader {
     ///     base of the checkpoint files, or a URL to an archive containing the checkpoint files.
     ///   - modelName: A distinct name for the model, to ensure that checkpoints with the same base 
     ///     name but for different models don't collide when downloaded.
-    public init(checkpointLocation: URL, modelName: String, additionalFiles: [String] = []) throws {
+    public init(
+        checkpointLocation: URL, modelName: String, additionalFiles: [String] = [],
+        fileSystem: FileSystem = FoundationFileSystem()
+    ) throws {
+        self.fileSystem = fileSystem
         let temporaryDirectory = FileManager.default.temporaryDirectory.appendingPathComponent(
             modelName, isDirectory: true)
 
@@ -72,20 +78,23 @@ open class CheckpointReader {
         if finalCheckpointLocation.isFileURL {
             self.localCheckpointLocation = finalCheckpointLocation
             indexReader = try CheckpointIndexReader(
-                file: finalCheckpointLocation.appendingPathExtension("index"))
+                file: finalCheckpointLocation.appendingPathExtension("index"),
+                fileSystem: fileSystem)
             self.header = try indexReader.readHeader()
         } else {
             let temporaryCheckpointBase = temporaryDirectory.appendingPathComponent(checkpointBase)
             self.localCheckpointLocation = temporaryCheckpointBase
             let localIndexFileLocation = temporaryCheckpointBase.appendingPathExtension("index")
             if FileManager.default.fileExists(atPath: localIndexFileLocation.path) {
-                indexReader = try CheckpointIndexReader(file: localIndexFileLocation)
+                indexReader = try CheckpointIndexReader(file: localIndexFileLocation,
+                    fileSystem: fileSystem)
                 self.header = try indexReader.readHeader()
             } else {
                 // The index file contains the number of shards, so obtain that first.
                 try CheckpointReader.downloadIndexFile(
                     from: finalCheckpointLocation, to: temporaryDirectory)
-                indexReader = try CheckpointIndexReader(file: localIndexFileLocation)
+                indexReader = try CheckpointIndexReader(file: localIndexFileLocation,
+                    fileSystem: fileSystem)
                 self.header = try indexReader.readHeader()
 
                 try CheckpointReader.downloadCheckpointFiles(
@@ -269,10 +278,9 @@ open class CheckpointReader {
         } else {
             do {
                 // It is far too slow to read the shards in each time a tensor is accessed, so we
-                // read the entire shard into an in-memory cache on first access. A better approach
-                // to mapping these files may be needed, because .alwaysMapped doesn't seem to help
-                // as much as it should.
-                let shardBytes = try Data(contentsOf: file, options: .alwaysMapped)
+                // read the entire shard into an in-memory cache on first access.
+                let shardFile = fileSystem.open(file.path)
+                let shardBytes = try shardFile.read()
                 shardCache[file] = shardBytes
                 return shardBytes
             } catch {

Support/Checkpoints/CheckpointReader.swift Checkpoints/CheckpointReader.swift

@@ -13,21 +13,27 @@
 // limitations under the License.
 
 import Foundation
+import ModelSupport
 import TensorFlow
 
 /// A Swift-native TensorFlow v2 checkpoint writer. This writer has no dependencies
 /// on the TensorFlow runtime or libraries.
 open class CheckpointWriter {
     // TODO: Extend handling to different tensor types.
     let tensors: [String: Tensor<Float>]
+    let fileSystem: FileSystem
 
     /// Initializes the checkpoint reader from a dictionary of tensors, keyed on their string names.
     ///
     /// - Parameters:
     ///   - tensors: A dictionary containing the tensors to be written, with the keys being the
     ///     names of those tensors to write in the checkpoint.
-    public init(tensors: [String: Tensor<Float>]) {
+    ///   - fileSystem: The filesystem used for writing the checkpoint.
+    public init(
+      tensors: [String: Tensor<Float>], fileSystem: FileSystem = FoundationFileSystem()
+    ) {
         self.tensors = tensors
+        self.fileSystem = fileSystem
     }
 
     /// Writes the checkpoint to disk, in a specified directory. A TensorFlow v2 checkpoint consists
@@ -40,11 +46,12 @@ open class CheckpointWriter {
     ///   - name: The base name of the checkpoint, which is what will have the .index and
     ///     .data-0000X-of-0000Y extensions appended to it for files in the checkpoint directory.
     public func write(to directory: URL, name: String) throws {
-        try createDirectoryIfMissing(at: directory.path)
+        try fileSystem.createDirectoryIfMissing(at: directory.path)
         let indexWriter = CheckpointIndexWriter(tensors: tensors)
         let indexHeader = indexWriter.serializedHeader()
         let headerLocation = directory.appendingPathComponent("\(name).index")
-        try indexHeader.write(to: headerLocation)
+        let headerFile = fileSystem.open(headerLocation.path)
+        try headerFile.write(indexHeader)
 
         // TODO: Handle splitting into multiple shards.
         try writeShard(
@@ -72,6 +79,7 @@ open class CheckpointWriter {
             }
         }
 
-        try outputBuffer.write(to: shardFile)
+        let outputFile = fileSystem.open(shardFile.path)
+        try outputFile.write(outputBuffer)
     }
 }

Support/Checkpoints/CheckpointWriter.swift Checkpoints/CheckpointWriter.swift

@@ -23,6 +23,7 @@ import Foundation
 public enum SnappyDecompressionError: Error {
     case illegalLiteralLength(upperBits: UInt8)
     case impossibleTagType(tagType: UInt8)
+    case uncompressedDataLengthMismatch(target: Int, actual: Int)
 }
 
 // The following extension to Data provides methods that read variable-length byte sequences
@@ -150,9 +151,8 @@ public extension Data {
             }
         }
         if uncompressedData.count != uncompressedLength {
-            // TODO: Determine if this should be elevated to a thrown error.
-            printError(
-                "Warning: uncompressed data length of \(uncompressedData.count) did not match desired length of \(uncompressedLength).")
+            throw SnappyDecompressionError.uncompressedDataLengthMismatch(
+                target: uncompressedLength, actual: uncompressedData.count)
         }
 
         return uncompressedData

Support/Checkpoints/Protobufs/tensor_bundle.pb.swift Checkpoints/Protobufs/tensor_bundle.pb.swift

@@ -10,7 +10,7 @@ set_target_properties(FastStyleTransfer PROPERTIES
 target_compile_options(FastStyleTransfer PRIVATE
   $<$<BOOL:${BUILD_TESTING}>:-enable-testing>)
 target_link_libraries(FastStyleTransfer PUBLIC
-  ModelSupport)
+  Checkpoints)
 
 add_executable(FastStyleTransferDemo
   Demo/ColabDemo.ipynb

Support/Checkpoints/Protobufs/tensor_shape.pb.swift Checkpoints/Protobufs/tensor_shape.pb.swift

@@ -1,6 +1,6 @@
 import FastStyleTransfer
 import Foundation
-import ModelSupport
+import Checkpoints
 import TensorFlow
 
 extension TransformerNet: ImportableLayer {}

Support/Checkpoints/Protobufs/tensor_slice.pb.swift Checkpoints/Protobufs/tensor_slice.pb.swift

@@ -1,5 +1,5 @@
 import TensorFlow
-import ModelSupport
+import Checkpoints
 
 public protocol ImportableLayer: KeyPathIterable {}
 

Support/Checkpoints/Protobufs/types.pb.swift Checkpoints/Protobufs/types.pb.swift

@@ -27,7 +27,7 @@ set_target_properties(MiniGo PROPERTIES
 target_compile_options(MiniGo PRIVATE
   $<$<BOOL:${BUILD_TESTING}>:-enable-testing>)
 target_link_libraries(MiniGo PUBLIC
-  ModelSupport)
+  Checkpoints)
 
 add_executable(MiniGoDemo
   main.swift)

Support/Checkpoints/Protobufs/versions.pb.swift Checkpoints/Protobufs/versions.pb.swift

@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import Checkpoints
 import TensorFlow
-import ModelSupport
 
 public class MiniGoCheckpointReader: CheckpointReader {
     private var layerCounts: [String: Int] = [:]

Support/Checkpoints/SnappyDecompression.swift Checkpoints/SnappyDecompression.swift

@@ -12,10 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import Foundation
-import TensorFlow
+import Checkpoints
 import Datasets
+import Foundation
 import ModelSupport
+import TensorFlow
 
 /// Represents a type that can contribute to the regularization term when training models.
 public protocol Regularizable: Differentiable {

FastStyleTransfer/CMakeLists.txt

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import Checkpoints
 import Datasets
 import Foundation
 import ModelSupport

FastStyleTransfer/Demo/Helpers.swift

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import Checkpoints
 import Foundation
 import ModelSupport
 import TensorFlow

FastStyleTransfer/Utility/ImportableLayer.swift

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import Checkpoints
 import ModelSupport
 import TensorFlow
 

MiniGo/CMakeLists.txt

@@ -10,6 +10,7 @@ let package = Package(
     ],
     products: [
         .library(name: "Batcher", targets: ["Batcher"]),
+        .library(name: "Checkpoints", targets: ["Checkpoints"]),
         .library(name: "Datasets", targets: ["Datasets"]),
         .library(name: "ModelSupport", targets: ["ModelSupport"]),
         .library(name: "ImageClassificationModels", targets: ["ImageClassificationModels"]),
@@ -26,14 +27,17 @@ let package = Package(
     ],
     targets: [
         .target(name: "Batcher", path: "Batcher"),
+        .target(
+            name: "Checkpoints", dependencies: ["SwiftProtobuf", "ModelSupport"],
+            path: "Checkpoints"),
         .target(name: "Datasets", dependencies: ["ModelSupport"], path: "Datasets"),
         .target(name: "STBImage", path: "Support/STBImage"),
         .target(
             name: "ModelSupport", dependencies: ["SwiftProtobuf", "STBImage"], path: "Support",
             exclude: ["STBImage"]),
         .target(name: "ImageClassificationModels", path: "Models/ImageClassification"),
         .target(name: "VideoClassificationModels", path: "Models/Spatiotemporal"),
-        .target(name: "TextModels", dependencies: ["Datasets"], path: "Models/Text"),
+        .target(name: "TextModels", dependencies: ["Checkpoints", "Datasets"], path: "Models/Text"),
         .target(name: "RecommendationModels", path: "Models/Recommendation"),
         .target(name: "TrainingLoop", dependencies: ["ModelSupport"], path: "TrainingLoop"),
         .target(
@@ -76,7 +80,7 @@ let package = Package(
             dependencies: ["Datasets", "ImageClassificationModels", "TrainingLoop"],
             path: "Examples/MobileNetV2-Imagenette"),
         .target(
-            name: "MiniGo", dependencies: ["ModelSupport"], path: "MiniGo", exclude: ["main.swift"]),
+            name: "MiniGo", dependencies: ["Checkpoints"], path: "MiniGo", exclude: ["main.swift"]),
         .target(
             name: "MiniGoDemo", dependencies: ["MiniGo"], path: "MiniGo", sources: ["main.swift"]),
         .target(
@@ -100,7 +104,7 @@ let package = Package(
         .target(name: "GAN", dependencies: ["Datasets", "ModelSupport"], path: "GAN"),
         .target(name: "DCGAN", dependencies: ["Datasets", "ModelSupport"], path: "DCGAN"),
         .target(
-            name: "FastStyleTransfer", dependencies: ["ModelSupport"], path: "FastStyleTransfer",
+            name: "FastStyleTransfer", dependencies: ["Checkpoints"], path: "FastStyleTransfer",
             exclude: ["Demo"]),
         .target(
             name: "FastStyleTransferDemo", dependencies: ["FastStyleTransfer"],
@@ -113,7 +117,7 @@ let package = Package(
                 "TextModels"
             ],
             path: "Benchmarks"),
-        .testTarget(name: "CheckpointTests", dependencies: ["ModelSupport"]),
+        .testTarget(name: "CheckpointTests", dependencies: ["Checkpoints"]),
         .target(
             name: "BERT-CoLA", dependencies: ["TextModels", "Datasets"], path: "Examples/BERT-CoLA"),
         .testTarget(name: "SupportTests", dependencies: ["ModelSupport"]),

MiniGo/Models/MiniGoCheckpointReader.swift

@@ -2,17 +2,9 @@ add_subdirectory(STBImage)
 
 add_library(ModelSupport
   BijectiveDictionary.swift
-  Checkpoints/CheckpointIndexReader.swift
-  Checkpoints/CheckpointIndexWriter.swift
-  Checkpoints/CheckpointReader.swift
-  Checkpoints/CheckpointWriter.swift
-  Checkpoints/Protobufs/tensor_bundle.pb.swift
-  Checkpoints/Protobufs/tensor_shape.pb.swift
-  Checkpoints/Protobufs/tensor_slice.pb.swift
-  Checkpoints/Protobufs/types.pb.swift
-  Checkpoints/Protobufs/versions.pb.swift
-  Checkpoints/SnappyDecompression.swift
   FileManagement.swift
+  FileSystem.swift
+  FoundationFileSystem.swift
   LabeledData.swift
   Image.swift
   Stderr.swift

Models/Text/BERT.swift

@@ -0,0 +1,36 @@
+// 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
+
+public protocol FileSystem {
+  /// Creates a directory at a path, if missing. If the directory exists, this does nothing.
+  ///
+  /// - Parameters:
+  ///   - path: The path of the desired directory.
+  func createDirectoryIfMissing(at path: String) throws
+
+  /// Opens a file at the specified location for reading or writing.
+  ///
+  /// - Parameters:
+  ///   - path: The path of the file to be opened.
+  func open(_ path: String) -> File
+}
+
+public protocol File {
+  func read() throws -> Data
+  func read(position: Int, count: Int) throws -> Data
+  func write(_ value: Data) throws
+  func write(_ value: Data, position: Int) throws
+}