Convert NeuMF-MovieLens to Epoches (#605)

Author: xihui-wu

Datasets/MovieLens/MovieLens.swift

@@ -34,20 +34,29 @@ extension Sequence where Iterator.Element: Hashable {
     }
 }
 
-public struct MovieLens {
+public struct MovieLens<Entropy: RandomNumberGenerator> {
     public let trainUsers: [Float]
     public let testUsers: [Float]
     public let testData: [[Float]]
     public let items: [Float]
     public let numUsers: Int
-    public let numItems: Int
-    public let trainMatrix: [TensorPair<Int32, Float>]
+    public let numItems: Int  
     public let user2id: [Float: Int]
     public let id2user: [Int: Float]
     public let item2id: [Float: Int]
     public let id2item: [Int: Float]
     public let trainNegSampling: Tensor<Float>
 
+    public typealias Samples = [TensorPair<Int32, Float>]
+    public typealias Batches = Slices<Sampling<Samples, ArraySlice<Int>>>
+    public typealias BatchedTensorPair = TensorPair<Int32, Float>
+    public typealias Training = LazyMapSequence<
+        TrainingEpochs<Samples, Entropy>, 
+        LazyMapSequence<Batches, BatchedTensorPair>
+      >
+    public let trainMatrix: Samples
+    public let training: Training
+
     static func downloadMovieLensDatasetIfNotPresent() -> URL {
         let localURL = DatasetUtilities.defaultDirectory.appendingPathComponent(
             "MovieLens", isDirectory: true)
@@ -60,7 +69,9 @@ public struct MovieLens {
         return dataFolder
     }
 
-    public init() {
+    public init(
+            trainBatchSize: Int = 1024, 
+            entropy: Entropy) {
         let trainFiles = try! String(
             contentsOf: MovieLens.downloadMovieLensDatasetIfNotPresent().appendingPathComponent(
                 "u1.base"), encoding: .utf8)
@@ -127,7 +138,28 @@ public struct MovieLens {
         self.id2user = id2user
         self.item2id = item2id
         self.id2item = id2item
-        self.trainMatrix = dataset
         self.trainNegSampling = trainNegSampling
+
+        self.trainMatrix = dataset
+        self.training = TrainingEpochs(
+            samples: trainMatrix, 
+            batchSize: trainBatchSize, 
+            entropy: entropy
+        ).lazy.map { (batches: Batches) -> LazyMapSequence<Batches, BatchedTensorPair> in
+            batches.lazy.map {
+                TensorPair<Int32, Float> (
+                    first: Tensor<Int32>($0.map(\.first)),
+                    second: Tensor<Float>($0.map(\.second))
+                )
+            }
+        }
+    }
+}
+
+extension MovieLens where Entropy == SystemRandomNumberGenerator {
+    public init(trainBatchSize: Int = 1024) {
+        self.init(
+            trainBatchSize: trainBatchSize, 
+            entropy: SystemRandomNumberGenerator())
     }
 }

Examples/NeuMF-MovieLens/main.swift

@@ -12,18 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import Batcher
 import Datasets
 import Foundation
 import RecommendationModels
 import TensorFlow
 
-let dataset = MovieLens()
+let dataset = MovieLens(trainBatchSize: 1024)
 let numUsers = dataset.numUsers
 let numItems = dataset.numItems
 
-let batcher = Batcher(on: dataset.trainMatrix, batchSize: 1024, shuffle: true)
-
 let size: [Int] = [16, 32, 16, 8]
 let regs: [Float] = [0.0, 0.0, 0.0, 0.0]
 
@@ -48,12 +45,13 @@ for element in dataset.testData {
 print("Dataset acquired.")
 
 print("Starting training...")
-for epoch in 1...20 {
+let epochCount = 20
+for (epoch, epochBatches) in dataset.training.prefix(epochCount).enumerated() {
     var avgLoss: Float = 0.0
     Context.local.learningPhase = .training
-    for data in batcher.sequenced() {
-        let userId = data.first
-        let rating = data.second
+    for batch in epochBatches {
+        let userId = batch.first
+        let rating = batch.second
         let (loss, grad) = valueWithGradient(at: model) { model -> Tensor<Float> in
             let logits = model(userId)
             return sigmoidCrossEntropy(logits: logits, labels: rating)