Port VGG16 to use new layer init API

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 = VGG16(classCount: 10)
+var model = makeVGG16(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

@@ -13,12 +13,50 @@
 // limitations under the License.
 
 import TensorFlow
+import LayerInit
 
 // Original Paper:
 // "Very Deep Convolutional Networks for Large-Scale Image Recognition"
 // Karen Simonyan, Andrew Zisserman
 // https://arxiv.org/abs/1409.1556
 
+public typealias AutoVGGBlock = AutoSequencedDefinition<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)]
+    }
+
+    return AutoSequencedMany(layers: blocks)
+        .then(AutoMaxPool2D(poolSize: (2, 2), strides: (2, 2)))
+}
+
+// TODO(shadaj): oh no
+public typealias AutoVGG16Backbone = AutoSequencedDefinition<AutoSequencedDefinition<AutoSequencedDefinition<AutoSequencedDefinition<AutoVGGBlock, AutoVGGBlock>, AutoVGGBlock>, AutoVGGBlock>, AutoVGGBlock>
+public typealias AutoVGG16 = AutoSequencedDefinition<AutoSequencedDefinition<AutoSequencedDefinition<AutoSequencedDefinition<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 VGGBlock: Layer {
     var blocks: [Conv2D<Float>] = []
     var maxpool = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
@@ -40,34 +78,6 @@ public struct VGGBlock: Layer {
     }
 }
 
-public struct VGG16: Layer {
-    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>
-
-    public init(classCount: Int = 1000) {
-        layer1 = VGGBlock(featureCounts: (3, 64, 64, 64), blockCount: 2)
-        layer2 = VGGBlock(featureCounts: (64, 128, 128, 128), blockCount: 2)
-        layer3 = VGGBlock(featureCounts: (128, 256, 256, 256), blockCount: 3)
-        layer4 = VGGBlock(featureCounts: (256, 512, 512, 512), blockCount: 3)
-        layer5 = VGGBlock(featureCounts: (512, 512, 512, 512), blockCount: 3)
-        output = Dense(inputSize: 4096, outputSize: classCount)
-    }
-
-    @differentiable
-    public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
-        let backbone = input.sequenced(through: layer1, layer2, layer3, layer4, layer5)
-        return backbone.sequenced(through: flatten, dense1, dense2, output)
-    }
-}
-
 public struct VGG19: Layer {
     var layer1: VGGBlock
     var layer2: VGGBlock

Models/LayerInit/AutoPool.swift

@@ -42,3 +42,46 @@ public struct AutoAvgPool2D<Scalar>: AutoLayer where Scalar: TensorFlowFloatingP
         ), outputShape)
     }
 }
+
+public struct AutoMaxPool2D<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
+    let poolSize: (Int, Int)
+    let strides: (Int, Int)
+    let padding: Padding
+
+    public typealias InstanceType = MaxPool2D<Scalar>
+    public typealias InputShape = (Int, Int, Int)
+    public typealias OutputShape = (Int, Int, Int)
+
+    public init(
+        poolSize: (Int, Int),
+        strides: (Int, Int) = (1, 1),
+        padding: Padding = .valid
+    ) {
+        self.poolSize = poolSize
+        self.strides = strides
+        self.padding = padding
+    }
+
+    public func buildModelWithOutputShape(inputShape: (Int, Int, Int)) -> (InstanceType, (Int, Int, Int)) {
+        let outputShape: (Int, Int, Int)
+        if (padding == .valid) {
+            outputShape = (
+                Int(ceil(Float(inputShape.0 - poolSize.0 + 1) / Float(strides.0))),
+                Int(ceil(Float(inputShape.1 - poolSize.1 + 1) / Float(strides.1))),
+                inputShape.2
+            )
+        } else {
+            outputShape = (
+                Int(ceil(Float(inputShape.0) / Float(strides.0))),
+                Int(ceil(Float(inputShape.1) / Float(strides.1))),
+                inputShape.2
+            )
+        }
+        
+        return (MaxPool2D<Scalar>(
+            poolSize: poolSize,
+            strides: strides,
+            padding: padding
+        ), outputShape)
+    }
+}

Models/LayerInit/AutoSequenced.swift

@@ -27,3 +27,37 @@ extension AutoLayer {
         return AutoSequencedDefinition<Self, T>(first: self, second: other)
     }
 }
+
+public struct AutoSequencedManyInstance<LayerType: Layer>: Layer
+where LayerType.Input == LayerType.Output {
+    var layers: [LayerType]
+
+    @differentiable
+    public func callAsFunction(_ input: LayerType.Input) -> LayerType.Output {
+        return layers.differentiableReduce(input) { $1($0) }
+    }
+}
+
+public struct AutoSequencedMany<LayerType: AutoLayer>: AutoLayer
+where
+  LayerType.OutputShape == LayerType.InputShape,
+  LayerType.InstanceType.Input == LayerType.InstanceType.Output {
+    let layers: [LayerType]
+
+    public typealias InstanceType = AutoSequencedManyInstance<LayerType.InstanceType>
+
+    public init(layers: [LayerType]) {
+        self.layers = layers
+    }
+
+    public func buildModelWithOutputShape(inputShape: LayerType.InputShape) -> (InstanceType, LayerType.OutputShape) {
+        var lastOutputShape = inputShape
+        let builtInstances = self.layers.map({ autoLayer -> LayerType.InstanceType in
+            let (instance, outputShape) = autoLayer.buildModelWithOutputShape(inputShape: lastOutputShape)
+            lastOutputShape = outputShape
+            return instance
+        })
+
+        return (AutoSequencedManyInstance(layers: builtInstances), lastOutputShape)
+    }
+}

Package.swift

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