Migrating the COCO dataset to Epochs

Datasets/CMakeLists.txt

@@ -25,7 +25,6 @@ add_library(Datasets
   ImageSegmentationDataset.swift
   OxfordIIITPets/OxfordIIITPets.swift)
 target_link_libraries(Datasets PUBLIC
-  Batcher
   ModelSupport)
 set_target_properties(Datasets PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})

Datasets/COCO/COCO.swift

@@ -1,3 +1,18 @@
+// 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
 
 // Code below is ported from https://github.com/cocometadata/cocoapi

Datasets/COCO/COCODataset.swift

@@ -1,52 +1,102 @@
-import Batcher
+// 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
+import TensorFlow
+
+public struct COCODataset<Entropy: RandomNumberGenerator> {
+  /// Type of the collection of non-collated batches.
+  public typealias Batches = Slices<Sampling<[ObjectDetectionExample], ArraySlice<Int>>>
+  /// The type of the training data, represented as a sequence of epochs, which
+  /// are collection of batches.
+  public typealias Training = LazyMapSequence<
+    TrainingEpochs<[ObjectDetectionExample], Entropy>,
+    LazyMapSequence<Batches, [ObjectDetectionExample]>
+  >
+  /// The type of the validation data, represented as a collection of batches.
+  public typealias Validation = LazyMapSequence<Slices<[ObjectDetectionExample]>, [ObjectDetectionExample]>
+  /// The training epochs.
+  public let training: Training
+  /// The validation batches.
+  public let validation: Validation
+
+  /// Creates an instance with `batchSize` on `device` using `remoteBinaryArchiveLocation`.
+  ///
+  /// - Parameters:
+  ///   - training: The COCO metadata for the training data.
+  ///   - validation: The COCO metadata for the validation data.
+  ///   - includeMasks: Whether to include the segmentation masks when loading the dataset.
+  ///   - batchSize: Number of images provided per batch.
+  ///   - entropy: A source of randomness used to shuffle sample ordering.  It
+  ///     will be stored in `self`, so if it is only pseudorandom and has value
+  ///     semantics, the sequence of epochs is deterministic and not dependent
+  ///     on other operations.
+  ///   - device: The Device on which resulting Tensors from this dataset will be placed, as well
+  ///     as where the latter stages of any conversion calculations will be performed.
+  public init(
+    training: COCO, validation: COCO, includeMasks: Bool, batchSize: Int,
+    entropy: Entropy, device: Device,
+    transform: @escaping (ObjectDetectionExample) -> [ObjectDetectionExample]
+  ) {
+    let trainingSamples = loadCOCOExamples(
+      from: training,
+      includeMasks: includeMasks,
+      batchSize: batchSize)
 
-public struct COCODataset: ObjectDetectionDataset {
-    public typealias SourceDataSet = [ObjectDetectionExample]
-    public let trainingExamples: SourceDataSet
-    public let training: Batcher<SourceDataSet>
-    public let testExamples: SourceDataSet
-    public let test: Batcher<SourceDataSet>
+    self.training = TrainingEpochs(samples: trainingSamples, batchSize: batchSize, entropy: entropy)
+      .lazy.map { (batches: Batches) -> LazyMapSequence<Batches, [ObjectDetectionExample]> in
+        return batches.lazy.map {
+          makeBatch(samples: $0, device: device, transform: transform)
+        }
+      }
+
+    let validationSamples = loadCOCOExamples(
+      from: validation,
+      includeMasks: includeMasks,
+      batchSize: batchSize)
 
-    public init(
-        training: COCO, test: COCO,
-        includeMasks: Bool, batchSize: Int, numWorkers: Int
-    ) {
-        self.trainingExamples =
-            loadCOCOExamples(
-                from: training,
-                includeMasks: includeMasks,
-                batchSize: batchSize,
-                numWorkers: numWorkers)
-        self.training =
-            Batcher(
-                on: trainingExamples,
-                batchSize: batchSize,
-                numWorkers: numWorkers,
-                shuffle: true)
-        self.testExamples =
-            loadCOCOExamples(
-                from: test,
-                includeMasks: includeMasks,
-                batchSize: batchSize,
-                numWorkers: numWorkers)
-        self.test =
-            Batcher(
-                on: testExamples,
-                batchSize: batchSize,
-                numWorkers: numWorkers,
-                shuffle: false)
+    self.validation = validationSamples.inBatches(of: batchSize).lazy.map {
+      makeBatch(samples: $0, device: device, transform: transform)
     }
+  }
+
+  public static func identity(_ example: ObjectDetectionExample) -> [ObjectDetectionExample] {
+    return [example]
+  }
 }
 
-func loadCOCOExamples(from coco: COCO, includeMasks: Bool, batchSize: Int, numWorkers: Int)
+extension COCODataset: ObjectDetectionData where Entropy == SystemRandomNumberGenerator {
+  /// Creates an instance with `batchSize`, using the SystemRandomNumberGenerator.
+  public init(
+    training: COCO, validation: COCO, includeMasks: Bool, batchSize: Int,
+    on device: Device = Device.default,
+    transform: @escaping (ObjectDetectionExample) -> [ObjectDetectionExample] = COCODataset.identity
+  ) {
+    self.init(
+      training: training, validation: validation, includeMasks: includeMasks, batchSize: batchSize,
+      entropy: SystemRandomNumberGenerator(), device: device, transform: transform)
+  }
+}
+
+
+func loadCOCOExamples(from coco: COCO, includeMasks: Bool, batchSize: Int)
     -> [ObjectDetectionExample]
 {
     let images = coco.metadata["images"] as! [COCO.Image]
     let batchCount: Int = images.count / batchSize + 1
-    let n = min(numWorkers, batchCount)
     let batches = Array(0..<batchCount)
-    let examples: [[ObjectDetectionExample]] = batches._concurrentMap(nthreads: n) { batchIdx in
+    let examples: [[ObjectDetectionExample]] = batches.map { batchIdx in
         var examples: [ObjectDetectionExample] = []
         for i in 0..<batchSize {
             let idx = batchSize * batchIdx + i
@@ -118,3 +168,12 @@ func loadCOCOExample(coco: COCO, image: COCO.Image, includeMasks: Bool) -> Objec
     }
     return ObjectDetectionExample(image: img, objects: objects)
 }
+
+fileprivate func makeBatch<BatchSamples: Collection>(
+  samples: BatchSamples, device: Device,
+  transform: (ObjectDetectionExample) -> [ObjectDetectionExample]
+) -> [ObjectDetectionExample] where BatchSamples.Element == ObjectDetectionExample {
+  return samples.reduce([]) {
+    $0 + transform($1)
+  }
+}

Datasets/COCO/COCOVariant.swift

@@ -1,3 +1,17 @@
+// 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
 import ModelSupport
 

Datasets/LanguageModelDataset.swift

@@ -1,3 +1,17 @@
+// 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 dataset suitable for language modeling.

Datasets/ObjectDetectionDataset.swift

@@ -1,4 +1,17 @@
-import Batcher
+// 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
 import ModelSupport
 import TensorFlow
@@ -52,7 +65,7 @@ public struct LabeledObject {
     }
 }
 
-public struct ObjectDetectionExample: _Collatable, KeyPathIterable {
+public struct ObjectDetectionExample: KeyPathIterable {
     public let image: LazyImage
     public let objects: [LabeledObject]
 
@@ -62,10 +75,28 @@ public struct ObjectDetectionExample: _Collatable, KeyPathIterable {
     }
 }
 
-public protocol ObjectDetectionDataset {
-    associatedtype SourceDataSet: Collection
-    where SourceDataSet.Element == ObjectDetectionExample, SourceDataSet.Index == Int
+/// Types whose elements represent an object detection dataset (with both
+/// training and validation data).
+public protocol ObjectDetectionData {
+  /// The type of the training data, represented as a sequence of epochs, which
+  /// are collection of batches.
+  associatedtype Training: Sequence
+  where Training.Element: Collection, Training.Element.Element == [ObjectDetectionExample]
+  /// The type of the validation data, represented as a collection of batches.
+  associatedtype Validation: Collection where Validation.Element == [ObjectDetectionExample]
+  /// Creates an instance from a given `batchSize`.
+  init(
+    training: COCO, validation: COCO, includeMasks: Bool, batchSize: Int, on device: Device,
+    transform: @escaping (ObjectDetectionExample) -> [ObjectDetectionExample])
+  /// The `training` epochs.
+  var training: Training { get }
+  /// The `validation` batches.
+  var validation: Validation { get }
 
-    var training: Batcher<SourceDataSet> { get }
-    var test: Batcher<SourceDataSet> { get }
+  // The following is probably going to be necessary since we can't extract that
+  // information from `Epochs` or `Batches`.
+  /// The number of samples in the `training` set.
+  //var trainingSampleCount: Int {get}
+  /// The number of samples in the `validation` set.
+  //var validationSampleCount: Int {get}
 }

Datasets/TensorPair.swift

@@ -13,15 +13,14 @@
 // limitations under the License.
 
 import TensorFlow
-import Batcher
 
 /// A generic tuple of two tensors `Tensor`.
 /// 
 /// - Note: `TensorPair` has a generic name and provides little semantic information, to conform to
 /// `Collatable`. You can use it for most basic datasets with one tensor of inputs and one tensor of
 /// labels but you should write your own struct for more complex tasks (or if you want more descriptive
 /// names).
-public struct TensorPair<S1: TensorFlowScalar, S2: TensorFlowScalar>: _Collatable, KeyPathIterable {
+public struct TensorPair<S1: TensorFlowScalar, S2: TensorFlowScalar>: KeyPathIterable {
     public var first: Tensor<S1>
     public var second: Tensor<S2>
 
@@ -30,4 +29,4 @@ public struct TensorPair<S1: TensorFlowScalar, S2: TensorFlowScalar>: _Collatabl
         self.first = first
         self.second = second
     }
-}
+}

Package.swift

@@ -25,7 +25,7 @@ let package = Package(
     ],
     targets: [
         .target(name: "Batcher", path: "Batcher"),
-        .target(name: "Datasets", dependencies: ["ModelSupport", "Batcher"], path: "Datasets"),
+        .target(name: "Datasets", dependencies: ["ModelSupport"], path: "Datasets"),
         .target(name: "STBImage", path: "Support/STBImage"),
         .target(
             name: "ModelSupport", dependencies: ["SwiftProtobuf", "STBImage"], path: "Support",
@@ -117,7 +117,7 @@ let package = Package(
         ),
         .target(
             name: "pix2pix",
-            dependencies: ["Batcher", "ArgumentParser", "ModelSupport", "Datasets"],
+            dependencies: ["ArgumentParser", "ModelSupport", "Datasets"],
             path: "pix2pix"
         ),
         .target(

Tests/DatasetsTests/COCO/COCODatasetTests.swift

@@ -9,26 +9,33 @@ final class COCODatasetTests: XCTestCase {
         // to avoid fetching the full training data during CI runs.
         let dataset = COCODataset(
             training: COCOVariant.loadVal(),
-            test: COCOVariant.loadTest(),
-            includeMasks: false, batchSize: 32, numWorkers: 8)
-        verify(dataset.trainingExamples)
-        verify(dataset.testExamples)
+            validation: COCOVariant.loadTest(),
+            includeMasks: false, batchSize: 32)
+
+        for epochBatches in dataset.training.prefix(1) {
+            let batch = epochBatches.first!
+            XCTAssertTrue(batch[0].image.width != 0)
+        }
+
+        let validationBatch = dataset.validation.first!
+        XCTAssertTrue(validationBatch[0].image.width != 0)
     }
 
     func testExamplesIncludingMasks() {
         // We use val/test variants here, instead of train/val, 
         // to avoid fetching the full training data during CI runs.
         let dataset = COCODataset(
             training: COCOVariant.loadVal(),
-            test: COCOVariant.loadTest(),
-            includeMasks: true, batchSize: 32, numWorkers: 8)
-        verify(dataset.trainingExamples)
-        verify(dataset.testExamples)
-    }
+            validation: COCOVariant.loadTest(),
+            includeMasks: true, batchSize: 32)
+
+        for epochBatches in dataset.training.prefix(1) {
+            let batch = epochBatches.first!
+            XCTAssertTrue(batch[0].image.width != 0)
+        }
 
-    func verify(_ examples: [ObjectDetectionExample]) {
-        XCTAssertTrue(examples.count > 0)
-        XCTAssertTrue(examples[0].image.width != 0)
+        let validationBatch = dataset.validation.first!
+        XCTAssertTrue(validationBatch[0].image.width != 0)
     }
 
     static var allTests = [