[WIP]: Exploring removing <Float> everywhere.

Models/ImageClassification/LeNet-5.swift

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 import TensorFlow
+import ModelSupport
 
 // Original Paper:
 // "Gradient-Based Learning Applied to Document Recognition"
@@ -21,21 +22,22 @@ import TensorFlow
 //
 // Note: this implementation connects all the feature maps in the second convolutional layer.
 // Additionally, ReLU is used instead of sigmoid activations.
-
-public struct LeNet: Layer {
-    public var conv1 = Conv2D<Float>(filterShape: (5, 5, 1, 6), padding: .same, activation: relu)
-    public var pool1 = AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
-    public var conv2 = Conv2D<Float>(filterShape: (5, 5, 6, 16), activation: relu)
-    public var pool2 = AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
-    public var flatten = Flatten<Float>()
-    public var fc1 = Dense<Float>(inputSize: 400, outputSize: 120, activation: relu)
-    public var fc2 = Dense<Float>(inputSize: 120, outputSize: 84, activation: relu)
-    public var fc3 = Dense<Float>(inputSize: 84, outputSize: 10)
+public struct LeNet: FloatModel {
+    public var conv1 = Conv2D(filterShape: (5, 5, 1, 6), padding: .same, activation: relu)
+    public var pool1 = AvgPool2D(poolSize: (2, 2), strides: (2, 2))
+    public var conv2 = Conv2D(filterShape: (5, 5, 6, 16), activation: relu)
+    public var pool2 = AvgPool2D(poolSize: (2, 2), strides: (2, 2))
+    public var flatten = Flatten()
+    public var fc1 = Dense(inputSize: 400, outputSize: 120, activation: relu)
+    public var fc2 = Dense(inputSize: 120, outputSize: 84, activation: relu)
+    public var fc3 = Dense(inputSize: 84, outputSize: 10)
 
     public init() {}
 
+    public typealias Tensor = TensorF  // Lines below do not compile without this!
+
     @differentiable
-    public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
+    public func callAsFunction(_ input: Tensor) -> Tensor {
         let convolved = input.sequenced(through: conv1, pool1, conv2, pool2)
         return convolved.sequenced(through: flatten, fc1, fc2, fc3)
     }

Models/ImageClassification/VGG.swift

@@ -13,44 +13,45 @@
 // limitations under the License.
 
 import TensorFlow
+import ModelSupport
 
 // Original Paper:
 // "Very Deep Convolutional Networks for Large-Scale Image Recognition"
 // Karen Simonyan, Andrew Zisserman
 // https://arxiv.org/abs/1409.1556
 
-public struct VGGBlock: Layer {
-    var blocks: [Conv2D<Float>] = []
-    var maxpool = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
+public struct VGGBlock: FloatModel {
+    var blocks: [Self.Conv2D] = []
+    var maxpool = MaxPool2D(poolSize: (2, 2), strides: (2, 2))
 
     public init(featureCounts: (Int, Int, Int, Int), blockCount: Int) {
-        self.blocks = [Conv2D<Float>(filterShape: (3, 3, featureCounts.0, featureCounts.1),
+        self.blocks = [Conv2D(filterShape: (3, 3, featureCounts.0, featureCounts.1),
             padding: .same,
             activation: relu)]
         for _ in 1..<blockCount {
-            self.blocks += [Conv2D<Float>(filterShape: (3, 3, featureCounts.2, featureCounts.3),
+            self.blocks += [Conv2D(filterShape: (3, 3, featureCounts.2, featureCounts.3),
                 padding: .same,
                 activation: relu)]
         }
     }
 
     @differentiable
-    public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
+    public func callAsFunction(_ input: TensorF) -> TensorF {
         return maxpool(blocks.differentiableReduce(input) { $1($0) })
     }
 }
 
-public struct VGG16: Layer {
+public struct VGG16: FloatModel {
     var layer1: VGGBlock
     var layer2: VGGBlock
     var layer3: VGGBlock
     var layer4: VGGBlock
     var layer5: VGGBlock
 
-    var flatten = Flatten<Float>()
-    var dense1 = Dense<Float>(inputSize: 512 * 7 * 7, outputSize: 4096, activation: relu)
-    var dense2 = Dense<Float>(inputSize: 4096, outputSize: 4096, activation: relu)
-    var output: Dense<Float>
+    var flatten = Flatten()
+    var dense1 = Dense(inputSize: 512 * 7 * 7, outputSize: 4096, activation: relu)
+    var dense2 = Dense(inputSize: 4096, outputSize: 4096, activation: relu)
+    var output: Self.Dense
 
     public init(classCount: Int = 1000) {
         layer1 = VGGBlock(featureCounts: (3, 64, 64, 64), blockCount: 2)
@@ -68,17 +69,17 @@ public struct VGG16: Layer {
     }
 }
 
-public struct VGG19: Layer {
+public struct VGG19: FloatModel {
     var layer1: VGGBlock
     var layer2: VGGBlock
     var layer3: VGGBlock
     var layer4: VGGBlock
     var layer5: VGGBlock
 
-    var flatten = Flatten<Float>()
-    var dense1 = Dense<Float>(inputSize: 512 * 7 * 7, outputSize: 4096, activation: relu)
-    var dense2 = Dense<Float>(inputSize: 4096, outputSize: 4096, activation: relu)
-    var output: Dense<Float>
+    var flatten = Flatten()
+    var dense1 = Dense(inputSize: 512 * 7 * 7, outputSize: 4096, activation: relu)
+    var dense2 = Dense(inputSize: 4096, outputSize: 4096, activation: relu)
+    var output: Self.Dense
 
     public init(classCount: Int = 1000) {
         layer1 = VGGBlock(featureCounts: (3, 64, 64, 64), blockCount: 2)

Models/Recommendation/DLRM.swift

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 import TensorFlow
+import ModelSupport
 
 /// The DLRM model is parameterized to support multiple ways of combining the latent spaces of the inputs.
 public enum InteractionType {
@@ -38,11 +39,12 @@ public enum InteractionType {
 /// "Deep Learning Recommendation Model for Personalization and Recommendation Systems"
 /// Maxim Naumov et al.
 /// https://arxiv.org/pdf/1906.00091.pdf
-public struct DLRM: Module {
+public struct DLRM: TypedModule {
+    public typealias Scalar = Float
 
-    public var mlpBottom: MLP
-    public var mlpTop: MLP
-    public var latentFactors: [Embedding<Float>]
+    public var mlpBottom: Self.MLP
+    public var mlpTop: Self.MLP
+    public var latentFactors: [Self.Embedding]
     @noDerivative public let nDense: Int
     @noDerivative public let interaction: InteractionType
 
@@ -63,7 +65,7 @@ public struct DLRM: Module {
         mlpTop = MLP(dims: [topInput] + lnTop + [1], sigmoidLastLayer: true)
         latentFactors = lnEmb.map { embeddingSize -> Embedding<Float> in
             // Use a random uniform initialization to match the reference implementation.
-            let weights = Tensor<Float>(
+            let weights = Tensor(
                 randomUniform: [embeddingSize, mSpa],
                 lowerBound: Tensor(Float(-1.0)/Float(embeddingSize)),
                 upperBound: Tensor(Float(1.0)/Float(embeddingSize)))

Models/Recommendation/MLP.swift

@@ -13,10 +13,11 @@
 // limitations under the License.
 
 import TensorFlow
+import ModelSupport
 
 /// MLP is a multi-layer perceptron and is used as a component of the DLRM model
-public struct MLP: Layer {
-    public var blocks: [Dense<Float>] = []
+public struct MLP<Scalar: TensorFlowFloatingPoint>: Layer, TypedModule {
+    public var blocks: [Self.Dense] = []
 
     /// Randomly initializes a new multilayer perceptron from the given hyperparameters.
     ///
@@ -35,12 +36,14 @@ public struct MLP: Layer {
     }
 
     @differentiable
-    public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
+    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
         let blocksReduced = blocks.differentiableReduce(input) { last, layer in
             layer(last)
         }
         return blocksReduced
     }
-
 }
 
+extension TypedModule where Scalar: TensorFlowFloatingPoint {
+  public typealias MLP = RecommendationModels.MLP<Scalar>
+}

Package.swift

@@ -46,10 +46,10 @@ let package = Package(
         .target(
             name: "ModelSupport", dependencies: ["SwiftProtobuf", "STBImage"], path: "Support",
             exclude: ["STBImage"]),
-        .target(name: "ImageClassificationModels", path: "Models/ImageClassification"),
+        .target(name: "ImageClassificationModels", dependencies: ["ModelSupport"], path: "Models/ImageClassification"),
         .target(name: "VideoClassificationModels", path: "Models/Spatiotemporal"),
         .target(name: "TextModels", dependencies: ["Datasets"], path: "Models/Text"),
-        .target(name: "RecommendationModels", path: "Models/Recommendation"),
+        .target(name: "RecommendationModels", dependencies: ["ModelSupport"], path: "Models/Recommendation"),
         .target(
             name: "Autoencoder1D", dependencies: ["Datasets", "ModelSupport"],
             path: "Autoencoder/Autoencoder1D"),

Support/SugarLayer.swift

@@ -0,0 +1,37 @@
+// Copyright 2019 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 model composed of a variety of layers all using the same scalar type.
+public protocol TypedModule: Module {
+  associatedtype Scalar: TensorFlowScalar
+
+  typealias Tensor = TensorFlow.Tensor<Scalar>  // For some reason, this doesn't work!
+  typealias TensorF = TensorFlow.Tensor<Float>
+  typealias TTensor = TensorFlow.Tensor
+}
+
+public extension TypedModule where Scalar: TensorFlowFloatingPoint {
+  typealias Conv2D = TensorFlow.Conv2D<Scalar>
+  typealias AvgPool2D = TensorFlow.AvgPool2D<Scalar>
+  typealias MaxPool2D = TensorFlow.MaxPool2D<Scalar>
+  typealias Flatten = TensorFlow.Flatten<Scalar>
+  typealias Dense = TensorFlow.Dense<Scalar>
+  typealias Embedding = TensorFlow.Embedding<Scalar>
+}
+
+/// A model composed of a variety of layers, where each one is parameterized
+/// by `Float`.
+public protocol FloatModel: TypedModule, Layer where Scalar == Float {}