Support ResNet-style models with skip connections

Author: shadaj

Examples/ResNet-CIFAR10/main.swift

@@ -21,7 +21,8 @@ let batchSize = 10
 let dataset = CIFAR10(batchSize: batchSize)
 
 // Use the network sized for CIFAR-10
-var model = ResNet(classCount: 10, depth: .resNet56, downsamplingInFirstStage: false)
+var model = autoResNet(classCount: 10, depth: .resNet56, downsamplingInFirstStage: false)
+    .buildModel(inputShape: (32, 32, 3))
 
 // the classic ImageNet optimizer setting diverges on CIFAR-10
 // let optimizer = SGD(for: model, learningRate: 0.1, momentum: 0.9)
@@ -33,6 +34,7 @@ for (epoch, epochBatches) in dataset.training.prefix(10).enumerated() {
     Context.local.learningPhase = .training
     var trainingLossSum: Float = 0
     var trainingBatchCount = 0
+    let batchCount = epochBatches.count
     for batch in epochBatches {
         let (images, labels) = (batch.data, batch.label)
         let (loss, gradients) = valueWithGradient(at: model) { model -> Tensor<Float> in
@@ -42,6 +44,7 @@ for (epoch, epochBatches) in dataset.training.prefix(10).enumerated() {
         trainingLossSum += loss.scalarized()
         trainingBatchCount += 1
         optimizer.update(&model, along: gradients)
+        print(trainingLossSum / Float(trainingBatchCount), Float(trainingBatchCount) / Float(batchCount))
     }
 
     Context.local.learningPhase = .inference

Models/ImageClassification/ResNet.swift

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 import TensorFlow
+import LayerInit
 
 // Original Paper:
 // "Deep Residual Learning for Image Recognition"
@@ -24,6 +25,12 @@ import TensorFlow
 // The structure of this implementation was inspired by the Flax ResNet example:
 // https://github.com/google/flax/blob/master/examples/imagenet/models.py
 
+public typealias AutoConvBN = AutoSequencedDefinition<AutoBatchNorm<(Int, Int, Int), Float>, AutoConv2D<Float>>
+public func autoConvBN(filterShape: (Int, Int), outputChannels: Int, strides: (Int, Int) = (1, 1), padding: Padding = .valid) -> AutoConvBN {
+    return AutoBatchNorm<(Int, Int, Int), Float>(momentum: 0.9, epsilon: 1e-5)
+        .then(AutoConv2D<Float>(filterShape: filterShape, outputChannels: outputChannels, strides: strides, padding: padding, useBias: false))
+}
+
 public struct ConvBN: Layer {
     public var conv: Conv2D<Float>
     public var norm: BatchNorm<Float>
@@ -43,6 +50,59 @@ public struct ConvBN: Layer {
     }
 }
 
+// TODO(shadaj): OH NO
+public typealias ConvPlusResidual = AutoSplitMerge<AutoSequencedMany<AutoConvBN>, AutoSequencedDefinition<AutoSequencedMany<AutoSequencedDefinition<AutoConvBN, AutoFunction<Tensor<Float>, Tensor<Float>, AutoConv2D<Float>.OutputShape, AutoConv2D<Float>.OutputShape>>>, AutoConvBN>, Tensor<Float>, AutoBatchNorm<(Int, Int, Int), Float>.InputShape>
+public typealias AutoResidualBlock = AutoSequencedDefinition<ConvPlusResidual, AutoFunction<Tensor<Float>, Tensor<Float>, (Int, Int, Int), (Int, Int, Int)>>
+public func autoResidualBlock(inputFilters: Int, filters: Int, strides: (Int, Int), useLaterStride: Bool, isBasic: Bool) -> AutoResidualBlock {
+    let outFilters = filters * (isBasic ? 1 : 4)
+    let needsProjection = (inputFilters != outFilters) || (strides.0 != 1)
+
+    let projection = needsProjection
+        ? autoConvBN(filterShape: (1, 1), outputChannels: outFilters, strides: strides)
+        : autoConvBN(filterShape: (1, 1), outputChannels: 1)
+
+    let residual = AutoSequencedMany(layers: needsProjection ? [projection]: [])
+
+    var earlyConvs: [AutoConvBN] = []
+    let lastConv: AutoConvBN
+    if isBasic {
+        earlyConvs = [
+            (autoConvBN(
+                filterShape: (3, 3), outputChannels: filters, strides: strides, padding: .same)),
+        ]
+        lastConv = autoConvBN(filterShape: (3, 3), outputChannels: outFilters, padding: .same)
+    } else {
+        if useLaterStride {
+            // Configure for ResNet V1.5 (the more common implementation).
+            earlyConvs.append(autoConvBN(filterShape: (1, 1), outputChannels: filters))
+            earlyConvs.append(
+                autoConvBN(filterShape: (3, 3), outputChannels: filters, strides: strides, padding: .same))
+        } else {
+            // Configure for ResNet V1 (the paper implementation).
+            earlyConvs.append(
+                autoConvBN(filterShape: (1, 1), outputChannels: filters, strides: strides))
+            earlyConvs.append(autoConvBN(filterShape: (3, 3), outputChannels: filters, padding: .same))
+        }
+        lastConv = autoConvBN(filterShape: (1, 1), outputChannels: outFilters)
+    }
+
+    let earlyConvsWithRelu = earlyConvs.map({ (conv) in
+        conv.then(AutoFunction(fnShape: { $0 }, fn: { (prev: Tensor<Float>) in relu(prev) }))
+    })
+
+    let lastConvResult = AutoSequencedMany(layers: earlyConvsWithRelu).then(lastConv)
+
+
+    let convPlusResidual = AutoSplitMerge(
+        layer1: residual,
+        layer2: lastConvResult,
+        mergeOutputShape: { (l1, l2) in l1 }, mergeFn: SplitMergeFunctionWrapper({ $0 + $1 }))
+
+    let finalResult = convPlusResidual.then(AutoFunction<Tensor<Float>, Tensor<Float>, (Int, Int, Int), (Int, Int, Int)>(fnShape: { $0 }, fn: { (prev: Tensor<Float>) in relu(prev) }))
+
+    return finalResult
+}
+
 public struct ResidualBlock: Layer {
     public var projection: ConvBN
     @noDerivative public let needsProjection: Bool
@@ -103,6 +163,50 @@ public struct ResidualBlock: Layer {
     }
 }
 
+public typealias AutoResNet = AutoSequencedDefinition<AutoSequencedDefinition<AutoSequencedDefinition<AutoSequencedDefinition<AutoSequencedDefinition<AutoConvBN, AutoFunction<Tensor<Float>, Tensor<Float>, AutoConv2D<Float>.OutputShape, AutoMaxPool2D<Float>.InputShape>>, AutoMaxPool2D<Float>>, AutoSequencedMany<AutoResidualBlock>>, AutoGlobalAvgPool2D<Float>>, AutoDense<Float>>
+public func autoResNet(
+    classCount: Int, depth: ResNet.Depth, downsamplingInFirstStage: Bool = true,
+    useLaterStride: Bool = true
+) -> AutoResNet {
+    let initialLayer: AutoConvBN
+    let maxPool: AutoMaxPool2D<Float>
+
+    let inputFilters: Int
+
+    if downsamplingInFirstStage {
+        inputFilters = 64
+        initialLayer = autoConvBN(
+            filterShape: (7, 7), outputChannels: inputFilters, strides: (2, 2), padding: .same)
+        maxPool = AutoMaxPool2D(poolSize: (3, 3), strides: (2, 2), padding: .same)
+    } else {
+        inputFilters = 16
+        initialLayer = autoConvBN(
+            filterShape: (3, 3), outputChannels: inputFilters, padding: .same)
+        maxPool = AutoMaxPool2D(poolSize: (1, 1), strides: (1, 1))  // no-op
+    }
+
+    var residualBlocks: [AutoResidualBlock] = []
+    var lastInputFilterCount = inputFilters
+    for (blockSizeIndex, blockSize) in depth.layerBlockSizes.enumerated() {
+        for blockIndex in 0..<blockSize {
+            let strides = ((blockSizeIndex > 0) && (blockIndex == 0)) ? (2, 2) : (1, 1)
+            let filters = inputFilters * Int(pow(2.0, Double(blockSizeIndex)))
+            let residualBlock = autoResidualBlock(
+                inputFilters: lastInputFilterCount, filters: filters, strides: strides,
+                useLaterStride: useLaterStride, isBasic: depth.usesBasicBlocks)//.buildModel(inputShape: (1, 1, lastInputFilterCount))
+            lastInputFilterCount = filters * (depth.usesBasicBlocks ? 1 : 4)
+            residualBlocks.append(residualBlock)
+        }
+    }
+
+    return initialLayer
+        .then(AutoFunction(fnShape: { $0 }, fn: { (prev: Tensor<Float>) in relu(prev) }))
+        .then(maxPool)
+        .then(AutoSequencedMany(layers: residualBlocks))
+        .then(AutoGlobalAvgPool2D())
+        .then(AutoDense(outputSize: classCount))
+}
+
 /// An implementation of the ResNet v1 and v1.5 architectures, at various depths.
 public struct ResNet: Layer {
     public var initialLayer: ConvBN

Models/LayerInit/AutoBatchNorm.swift

@@ -0,0 +1,33 @@
+import TensorFlow
+
+public struct AutoBatchNorm<Shape, Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
+    let axis: Int
+    let momentum: Scalar
+    let epsilon: Scalar
+
+    public typealias InstanceType = BatchNorm<Scalar>
+    public typealias InputShape = Shape
+    public typealias OutputShape = Shape
+
+    public init(
+        axis: Int = -1,
+        momentum: Scalar = 0.99,
+        epsilon: Scalar = 0.001
+    ) {
+        self.axis = axis
+        self.momentum = momentum
+        self.epsilon = epsilon
+    }
+
+    public func buildModelWithOutputShape(inputShape: Shape) -> (InstanceType, Shape) {
+        let inputShapeArray: [Int]
+        if let inputShapeTuple = inputShape as? (Int, Int, Int) {
+            inputShapeArray = [inputShapeTuple.0, inputShapeTuple.1, inputShapeTuple.2]
+        } else {
+            fatalError("Could not extract out elements of shape")
+        }
+
+        let featureCount = inputShapeArray[(inputShapeArray.count + axis) % inputShapeArray.count]
+        return (BatchNorm<Scalar>(featureCount: featureCount, axis: axis, momentum: momentum, epsilon: epsilon), inputShape)
+    }
+}

Models/LayerInit/AutoDense.swift

@@ -17,3 +17,75 @@ public struct AutoDense<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint
         return (Dense<Scalar>(inputSize: inputShape, outputSize: self.outputSize, activation: self.activation), self.outputSize)
     }
 }
+
+// Workaround https://bugs.swift.org/browse/TF-1122
+public final class SplitMergeFunctionWrapper<Output1: Differentiable, Output2: Differentiable, CommonOutput: Differentiable> {
+  public typealias F = @differentiable (Output1, Output2) -> CommonOutput 
+  public var f: F
+  public init(_ f: @escaping F) { self.f = f }
+}
+
+public struct SplitMergeInstance<Layer1: Layer, Layer2: Layer, CommonOutput: Differentiable>: Layer
+where Layer1.Input == Layer2.Input, Layer1.TangentVector.VectorSpaceScalar == Layer2.TangentVector.VectorSpaceScalar {
+    var layer1: Layer1
+    var layer2: Layer2
+    @noDerivative let mergeFn: SplitMergeFunctionWrapper<Layer1.Output, Layer2.Output, CommonOutput>
+    
+    public init(layer1: Layer1, layer2: Layer2, mergeFn: SplitMergeFunctionWrapper<Layer1.Output, Layer2.Output, CommonOutput>) {
+        self.layer1 = layer1
+        self.layer2 = layer2
+        self.mergeFn = mergeFn
+    }
+
+    @differentiable
+    public func callAsFunction(_ input: Layer1.Input) -> CommonOutput {
+        let layer1Out = layer1(input)
+        let layer2Out = layer2(input)
+        return mergeFn.f(layer1Out, layer2Out)
+    }
+}
+
+public struct AutoSplitMerge<Layer1: AutoLayer, Layer2: AutoLayer, CommonOutput: Differentiable, OutputShape>: AutoLayer
+where Layer1.InputShape == Layer2.InputShape, Layer1.InstanceType.Input == Layer2.InstanceType.Input, Layer1.InstanceType.TangentVector.VectorSpaceScalar == Layer2.InstanceType.TangentVector.VectorSpaceScalar {
+    let layer1: Layer1
+    let layer2: Layer2
+    
+    let mergeOutputShape: (Layer1.OutputShape, Layer2.OutputShape) -> OutputShape
+    let mergeFn: SplitMergeFunctionWrapper<Layer1.InstanceType.Output, Layer2.InstanceType.Output, CommonOutput>
+
+    public typealias InstanceType = SplitMergeInstance<Layer1.InstanceType, Layer2.InstanceType, CommonOutput>
+    public typealias InputShape = Layer1.InputShape
+    public typealias OutputShape = OutputShape
+
+    public init(layer1: Layer1, layer2: Layer2, mergeOutputShape: @escaping (Layer1.OutputShape, Layer2.OutputShape) -> OutputShape, mergeFn: SplitMergeFunctionWrapper<Layer1.InstanceType.Output, Layer2.InstanceType.Output, CommonOutput>) {
+        self.layer1 = layer1
+        self.layer2 = layer2
+        self.mergeOutputShape = mergeOutputShape
+        self.mergeFn = mergeFn
+    }
+
+    public func buildModelWithOutputShape(inputShape: Layer1.InputShape) -> (InstanceType, OutputShape) {
+        let (layer1Built, layer1OutputShape) = layer1.buildModelWithOutputShape(inputShape: inputShape)
+        let (layer2Built, layer2OutputShape) = layer2.buildModelWithOutputShape(inputShape: inputShape)
+        return (SplitMergeInstance(layer1: layer1Built, layer2: layer2Built, mergeFn: self.mergeFn), self.mergeOutputShape(layer1OutputShape, layer2OutputShape))
+    }
+}
+
+
+public struct AutoFunction<Input: Differentiable, Output: Differentiable, InputShape, OutputShape>: AutoLayer {
+    let fnShape: (InputShape) -> OutputShape
+    let fn: @differentiable (Input) -> Output
+
+    public typealias InstanceType = Function<Input, Output>
+    public typealias InputShape = InputShape
+    public typealias OutputShape = OutputShape
+
+    public init(fnShape: @escaping (InputShape) -> OutputShape, fn: @escaping @differentiable (Input) -> Output) {
+        self.fnShape = fnShape
+        self.fn = fn
+    }
+
+    public func buildModelWithOutputShape(inputShape: InputShape) -> (InstanceType, OutputShape) {
+        return (Function(fn), fnShape(inputShape))
+    }
+}

Models/LayerInit/AutoPool.swift

@@ -43,6 +43,19 @@ public struct AutoAvgPool2D<Scalar>: AutoLayer where Scalar: TensorFlowFloatingP
     }
 }
 
+public struct AutoGlobalAvgPool2D<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
+    public typealias InstanceType = GlobalAvgPool2D<Scalar>
+    public typealias InputShape = (Int, Int, Int)
+    public typealias OutputShape = Int
+
+    public init() {
+    }
+
+    public func buildModelWithOutputShape(inputShape: (Int, Int, Int)) -> (InstanceType, Int) {
+        return (GlobalAvgPool2D<Scalar>(), inputShape.2)
+    }
+}
+
 public struct AutoMaxPool2D<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
     let poolSize: (Int, Int)
     let strides: (Int, Int)

Models/LayerInit/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_library(LayerInit
+  AutoBatchNorm.swift
   AutoConv.swift
   AutoDense.swift
   AutoFlatten.swift