Explore LayerModule protocol to enable better type inference

Author: shadaj

Examples/VGG-Imagewoof/main.swift

@@ -19,7 +19,7 @@ import TensorFlow
 let batchSize = 32
 
 let dataset = Imagewoof(batchSize: batchSize, inputSize: .full, outputSize: 224)
-var model = makeVGG16(classCount: 10).buildModel(inputShape: (224, 224, 3))
+var model = AutoVGG16(classCount: 10).buildModel(inputShape: (224, 224, 3))
 let optimizer = SGD(for: model, learningRate: 0.02, momentum: 0.9, decay: 0.0005)
 
 print("Starting training...")

Models/ImageClassification/VGG.swift

@@ -20,41 +20,46 @@ import LayerInit
 // Karen Simonyan, Andrew Zisserman
 // https://arxiv.org/abs/1409.1556
 
-public typealias AutoVGGBlock = AutoSequenced<AutoSequencedMany<AutoConv2D<Float>>, AutoMaxPool2D<Float>>
-func makeVGGBlock(featureCounts: (Int, Int, Int, Int), blockCount: Int) -> AutoVGGBlock {
-    var blocks: [AutoConv2D<Float>] = [
-        AutoConv2D<Float>(filterShape: (3, 3), outputChannels: featureCounts.1,
-            padding: .same,
-            activation: relu)]
-    for _ in 1..<blockCount {
-        blocks += [AutoConv2D(filterShape: (3, 3), outputChannels: featureCounts.3,
-            padding: .same,
-            activation: relu)]
-    }
+public struct AutoVGGBlock: AutoModule {
+    let featureCounts: (Int, Int, Int, Int)
+    let blockCount: Int
 
-    return AutoSequencedMany(layers: blocks)
-        .then(AutoMaxPool2D(poolSize: (2, 2), strides: (2, 2)))
+    public typealias LayerType = AutoSequenced<AutoSequencedMany<AutoConv2D<Float>>, AutoMaxPool2D<Float>>
+    public lazy var initializeLayer: LayerType = {
+        var blocks: [AutoConv2D<Float>] = [
+            AutoConv2D<Float>(filterShape: (3, 3), outputChannels: featureCounts.1,
+                padding: .same,
+                activation: relu)]
+
+        for _ in 1..<blockCount {
+            blocks += [AutoConv2D(filterShape: (3, 3), outputChannels: featureCounts.3,
+                padding: .same,
+                activation: relu)]
+        }
+
+        return AutoSequencedMany(layers: blocks)
+            .then(AutoMaxPool2D(poolSize: (2, 2), strides: (2, 2)))
+    }()
 }
 
-// TODO(shadaj): oh no
-public typealias AutoVGG16Backbone = AutoSequenced<AutoSequenced<AutoSequenced<AutoSequenced<AutoVGGBlock, AutoVGGBlock>, AutoVGGBlock>, AutoVGGBlock>, AutoVGGBlock>
-public typealias AutoVGG16 = AutoSequenced<AutoSequenced<AutoSequenced<AutoSequenced<AutoVGG16Backbone, AutoFlatten<Float>>, AutoDense<Float>>, AutoDense<Float>>, AutoDense<Float>>
-
-public func makeVGG16(classCount: Int = 1000) -> AutoVGG16 {
-    let layer1 = makeVGGBlock(featureCounts: (3, 64, 64, 64), blockCount: 2)
-    let layer2 = makeVGGBlock(featureCounts: (64, 128, 128, 128), blockCount: 2)
-    let layer3 = makeVGGBlock(featureCounts: (128, 256, 256, 256), blockCount: 3)
-    let layer4 = makeVGGBlock(featureCounts: (256, 512, 512, 512), blockCount: 3)
-    let layer5 = makeVGGBlock(featureCounts: (512, 512, 512, 512), blockCount: 3)
-
-    let flatten = AutoFlatten<Float>()
-    let dense1 = AutoDense<Float>(outputSize: 4096, activation: relu)
-    let dense2 = AutoDense<Float>(outputSize: 4096, activation: relu)
-    let output = AutoDense<Float>(outputSize: classCount)
-
-    let backbone = layer1.then(layer2).then(layer3).then(layer4).then(layer5)
-    let fullModel = backbone.then(flatten).then(dense1).then(dense2).then(output)
-    return fullModel
+public struct AutoVGG16: AutoModule {
+    let classCount: Int
+
+    public init(classCount: Int = 1000) {
+        self.classCount = classCount
+    }
+
+    public lazy var initializeLayer = {
+        return AutoVGGBlock(featureCounts: (3, 64, 64, 64), blockCount: 2)
+            .then(AutoVGGBlock(featureCounts: (64, 128, 128, 128), blockCount: 2))
+            .then(AutoVGGBlock(featureCounts: (128, 256, 256, 256), blockCount: 3))
+            .then(AutoVGGBlock(featureCounts: (256, 512, 512, 512), blockCount: 3))
+            .then(AutoVGGBlock(featureCounts: (512, 512, 512, 512), blockCount: 3))
+            .then(AutoFlatten())
+            .then(AutoDense(outputSize: 4096, activation: relu))
+            .then(AutoDense(outputSize: 4096, activation: relu))
+            .then(AutoDense(outputSize: classCount))
+    }()
 }
 
 public struct VGGBlock: Layer {

Models/LayerInit/AutoModule.swift

@@ -0,0 +1,18 @@
+import TensorFlow
+
+public protocol AutoModule: AutoLayer {
+    associatedtype LayerType: AutoLayer
+
+    var initializeLayer: LayerType { mutating get }
+}
+
+extension AutoModule  {
+    public typealias InstanceType = LayerType.InstanceType
+    public typealias InputShape = LayerType.InputShape
+    public typealias OutputShape = LayerType.OutputShape
+
+    public func buildModelWithOutputShape(inputShape: InputShape) -> (InstanceType, OutputShape) {
+        var selfCopy = self
+        return selfCopy.initializeLayer.buildModelWithOutputShape(inputShape: inputShape)
+    }
+}