Introduce new layer initialization APIs with automatic shape computation

Models/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(ImageClassification)
+add_subdirectory(LayerInit)
 add_subdirectory(Recommendation)
 add_subdirectory(Text)

Models/LayerInit/AutoBatchNorm.swift

@@ -0,0 +1,28 @@
+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<Prefix>(inputShape: Shape, keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, Shape) {
+        let inputShapeArray: [Int] = intTupleToArray(tuple: inputShape)
+
+        let featureCount = inputShapeArray[(inputShapeArray.count + axis) % inputShapeArray.count]
+        return (BatchNorm<Scalar>(featureCount: featureCount, axis: axis, momentum: momentum, epsilon: epsilon), inputShape)
+    }
+}

Models/LayerInit/AutoConv.swift

@@ -0,0 +1,63 @@
+import TensorFlow
+
+public struct AutoConv2D<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
+    let filterShape: (Int, Int)
+    let outputChannels: Int
+    let strides: (Int, Int)
+    let padding: Padding
+    let dilations: (Int, Int)
+    let activation: Conv2D<Scalar>.Activation
+    let useBias: Bool
+    let filterInitializer: ParameterInitializer<Scalar>
+    let biasInitializer: ParameterInitializer<Scalar>
+
+    public typealias InstanceType = Conv2D<Scalar>
+    public typealias InputShape = (Int, Int, Int)
+    public typealias OutputShape = (Int, Int, Int)
+
+    public init(
+        filterShape: (Int, Int),
+        outputChannels: Int,
+        strides: (Int, Int) = (1, 1),
+        padding: Padding = .valid,
+        dilations: (Int, Int) = (1, 1),
+        activation: @escaping Conv2D<Scalar>.Activation = identity,
+        useBias: Bool = true,
+        filterInitializer: @escaping ParameterInitializer<Scalar> = glorotUniform(),
+        biasInitializer: @escaping ParameterInitializer<Scalar> = zeros()
+    ) {
+        self.filterShape = filterShape
+        self.outputChannels = outputChannels
+        self.strides = strides
+        self.padding = padding
+        self.dilations = dilations
+        self.activation = activation
+        self.useBias = useBias
+        self.filterInitializer = filterInitializer
+        self.biasInitializer = biasInitializer
+    }
+
+    public func buildModelWithOutputShape<Prefix>(inputShape: (Int, Int, Int), keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, (Int, Int, Int)) {
+        let outputShape: (Int, Int, Int)
+        if (padding == .valid) {
+            outputShape = (
+                Int(ceil(Float(inputShape.0 - filterShape.0 + 1) / Float(strides.0))),
+                Int(ceil(Float(inputShape.1 - filterShape.1 + 1) / Float(strides.1))),
+                outputChannels
+            )
+        } else {
+            outputShape = (
+                Int(ceil(Float(inputShape.0) / Float(strides.0))),
+                Int(ceil(Float(inputShape.1) / Float(strides.1))),
+                outputChannels
+            )
+        }
+
+        return (Conv2D<Scalar>(
+            filterShape: (filterShape.0, filterShape.1, inputShape.2, outputChannels),
+            strides: strides, padding: padding, dilations: dilations,
+            activation: activation, useBias: useBias,
+            filterInitializer: filterInitializer, biasInitializer: biasInitializer
+        ), outputShape)
+    }
+}

Models/LayerInit/AutoDense.swift

@@ -0,0 +1,19 @@
+import TensorFlow
+
+public struct AutoDense<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
+    let outputSize: Int;
+    let activation: Dense<Scalar>.Activation
+
+    public typealias InstanceType = Dense<Scalar>
+    public typealias InputShape = Int
+    public typealias OutputShape = Int
+
+    public init(outputSize: Int, activation: @escaping Dense<Scalar>.Activation = identity) {
+        self.outputSize = outputSize
+        self.activation = activation
+    }
+
+    public func buildModelWithOutputShape<Prefix>(inputShape: Int, keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, Int) {
+        return (Dense<Scalar>(inputSize: inputShape, outputSize: self.outputSize, activation: self.activation), self.outputSize)
+    }
+}

Models/LayerInit/AutoFlatten.swift

@@ -0,0 +1,13 @@
+import TensorFlow
+
+public struct AutoFlatten<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
+    public typealias InstanceType = Flatten<Scalar>
+    public typealias InputShape = Any
+    public typealias OutputShape = Int
+
+    public init() {}
+
+    public func buildModelWithOutputShape<Prefix>(inputShape: Any, keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, Int) {
+        return (Flatten<Scalar>(), intTupleToArray(tuple: inputShape).reduce(1, { $0 * $1 }))
+    }
+}

Models/LayerInit/AutoFunction.swift

@@ -0,0 +1,24 @@
+import TensorFlow
+
+/// A layer that applies a user-defined function to the input data
+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
+
+    /// Constructs a function layer instance.
+    /// Parameters:
+    ///     - fnShape: a function that computes the output shape of the function given the input shape
+    ///     - fn: a function that computes the output data of the function given the input data
+    public init(fnShape: @escaping (InputShape) -> OutputShape, fn: @escaping @differentiable (Input) -> Output) {
+        self.fnShape = fnShape
+        self.fn = fn
+    }
+
+    public func buildModelWithOutputShape<Prefix>(inputShape: InputShape, keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, OutputShape) {
+        return (Function(fn), fnShape(inputShape))
+    }
+}

Models/LayerInit/AutoLayer.swift

@@ -0,0 +1,53 @@
+import TensorFlow
+
+/// A layer "blueprint", which defines elements that can be constructed into a Layer instance for training
+public protocol AutoLayer {
+    /// The type of the layer instance that can be built with this prototype
+    associatedtype InstanceType: Layer
+
+    /// The specific tuple of `Int`s that define the input shape of the layer
+    associatedtype InputShape
+
+    /// The specific tuple of `Int`s that define the output shape of the layer
+    associatedtype OutputShape
+
+    /// Initializes a new instance of the layer defined by this blueprint.
+    /// Parameters:
+    ///     - inputShape: the shape of a single input instance (no batch) to this layer
+    ///     - keyPathSoFar: a `KeyPath` that tracks the path from the root layer to the current layer instance
+    ///     - keyDict: a dictionary tracking the mapping from `AutoLayerKey`s to the key path to the layer instance
+    /// Returns:
+    ///     - $0: the instance of the layer with the given input shape
+    ///     - $1: the output shape of the layer instance computed based on the input shape
+    func buildModelWithOutputShape<Prefix>(inputShape: InputShape, keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, OutputShape)
+}
+
+extension AutoLayer {
+    /// Builds an instance of the model with the given input shape
+    public func buildModel(inputShape: InputShape) -> BuiltAutoLayer<InstanceType> {
+        var keyDict: [AnyAutoLayerKey: Any] = [:]
+        let (layerInstance, _) = self.buildModelWithOutputShape(inputShape: inputShape, keyPathSoFar: \InstanceType.self, keyDict: &keyDict)
+        return BuiltAutoLayer(layer: layerInstance, keyMapping: keyDict)
+    }
+}
+
+/// A layer instance containing a model built with the AutoLayer API. Offers keyed access to layers with `AutoLayerKey`.
+public struct BuiltAutoLayer<InstanceType: Layer>: Layer {
+    public var layer: InstanceType
+    @noDerivative let keyMapping: [AnyAutoLayerKey: Any]
+
+    public init(layer: InstanceType, keyMapping: [AnyAutoLayerKey: Any]) {
+        self.layer = layer
+        self.keyMapping = keyMapping
+    }
+
+    @differentiable
+    public func callAsFunction(_ input: InstanceType.Input) -> InstanceType.Output {
+        return layer(input)
+    }
+
+    /// Grab a specific layer by the given `AutoLayerKey`.
+    public subscript<T>(index: AutoLayerKey<T>) -> T {
+        return self.layer[keyPath: self.keyMapping[index] as! KeyPath<InstanceType, T>]
+    }
+}

Models/LayerInit/AutoLayerKey.swift

@@ -0,0 +1,44 @@
+import TensorFlow
+
+public class AnyAutoLayerKey: Hashable {
+    public func hash(into hasher: inout Hasher) {
+        hasher.combine(ObjectIdentifier(self))
+    }
+
+    public static func == (lhs: AnyAutoLayerKey, rhs: AnyAutoLayerKey) -> Bool {
+        return lhs === rhs
+    }
+}
+
+/// A key that can be associated with an `AutoLayer` to access it after it has been built as part of a larger model.
+public class AutoLayerKey<T: Layer>: AnyAutoLayerKey {
+    public override init() {}
+}
+
+/// A layer blueprint that associates an underlying blueprint with an `AutoLayerKey` so that the underlying instance can be accessed from the built model.
+public struct KeyedAutoLayer<Underlying: AutoLayer>: AutoLayer {
+    let underlying: Underlying
+    let key: AutoLayerKey<InstanceType>
+
+    public typealias InstanceType = Underlying.InstanceType
+    public typealias InputShape = Underlying.InputShape
+    public typealias OutputShape = Underlying.OutputShape
+
+    public init(_ underlying: Underlying, key: AutoLayerKey<InstanceType>) {
+        self.underlying = underlying
+        self.key = key
+    }
+
+    public func buildModelWithOutputShape<Prefix>(inputShape: InputShape, keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, OutputShape) {
+        let (layer, outputShape) = underlying.buildModelWithOutputShape(inputShape: inputShape, keyPathSoFar: keyPathSoFar, keyDict: &keyDict)
+        keyDict[self.key] = keyPathSoFar
+        return (layer, outputShape)
+    }
+}
+
+extension AutoLayer {
+    /// Attaches a key to an existing layer blueprint
+    public func withKey(_ key: AutoLayerKey<InstanceType>) -> KeyedAutoLayer<Self> {
+        return KeyedAutoLayer(self, key: key)
+    }
+}

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<Prefix>(inputShape: InputShape, keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, OutputShape) {
+        var selfCopy = self
+        return selfCopy.initializeLayer.buildModelWithOutputShape(inputShape: inputShape, keyPathSoFar: keyPathSoFar, keyDict: &keyDict)
+    }
+}

Models/LayerInit/AutoPool.swift

@@ -0,0 +1,100 @@
+import TensorFlow
+
+public struct AutoAvgPool2D<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
+    let poolSize: (Int, Int)
+    let strides: (Int, Int)
+    let padding: Padding
+
+    public typealias InstanceType = AvgPool2D<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<Prefix>(inputShape: (Int, Int, Int), keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (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 (AvgPool2D<Scalar>(
+            poolSize: poolSize,
+            strides: strides,
+            padding: padding
+        ), outputShape)
+    }
+}
+
+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<Prefix>(inputShape: (Int, Int, Int), keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, Int) {
+        return (GlobalAvgPool2D<Scalar>(), inputShape.2)
+    }
+}
+
+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<Prefix>(inputShape: (Int, Int, Int), keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (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)
+    }
+}