Merge branch 'master' into complex-numbers2

Author: bartchr808

Sources/DeepLearning/Core/Dataset.swift

@@ -0,0 +1,220 @@
+#if !COMPILING_TENSORFLOW_MODULE
+import TensorFlow
+#endif
+
+//===-- Dataset.swift -----------------------------------------*- swift -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// The dataset API.
+//
+//===----------------------------------------------------------------------===//
+
+/// The default graph seed.
+///
+/// - Note: See TensorFlow's `python.framework.random_seed.DEFAULT_GRAPH_SEED`.
+@usableFromInline let _defaultGraphSeed: Int64 = 87654321
+
+/// Returns the local seeds an operation should use given an op-specific seed.
+///
+/// Given operation-specific seed, `seed`, this helper function returns two
+/// seeds derived from graph-level and op-level seeds. Many random operations
+/// internally use the two seeds to allow user to change the seed globally for a
+/// graph, or for only specific operations.
+///
+/// - Note: See TensorFlow's `python.framework.random_seed.get_seed`.
+///
+// TODO: There's no support for TF's "global seed" yet, so we always use the
+// default graph seed as the first seed. Need to investigate the best way to
+// model TF's "global seed".
+@usableFromInline @inline(__always)
+func _tensorSeeds(_ seed: Tensor<Int64>) -> (Tensor<Int64>, Tensor<Int64>) {
+  return (Tensor(_defaultGraphSeed), seed)
+}
+
+//===----------------------------------------------------------------------===//
+// Single value dataset
+//===----------------------------------------------------------------------===//
+
+/// Represents a potentially large set of elements.
+///
+/// A `Dataset` can be used to represent an input pipeline as a collection of
+/// element tensors.
+@_fixed_layout
+public struct Dataset<Element : TensorGroup> {
+  public let _handle: VariantHandle
+
+  @inlinable
+  public init(_handle: VariantHandle) {
+    self._handle = _handle
+  }
+}
+
+public extension Dataset {
+  @inlinable
+  init(randomSeed: Int64) {
+    let (seed1, seed2) = _tensorSeeds(Tensor(randomSeed))
+    self.init(_handle: Raw.experimentalRandomDataset(
+      seed: seed1,
+      seed2: seed2,
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList))
+  }
+}
+
+public extension Dataset {
+  /// Creates a dataset from a batch of elements as a tensor.
+  @inlinable
+  init(elements: Element) {
+    self.init(_handle: Raw.tensorSliceDataset(
+      components: [elements],
+      outputShapes: Element._unknownShapeList))
+  }
+}
+
+extension Dataset : Sequence {
+  public typealias Iterator = DatasetIterator<Element>
+
+  /// Returns an iterator over the elements of this dataset.
+  @inlinable
+  public func makeIterator() -> DatasetIterator<Element> {
+    let resource = Raw.anonymousIterator(
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList)
+    Raw.makeIterator(dataset: _handle, iterator: resource)
+    return DatasetIterator(_handle: resource)
+  }
+}
+
+public extension Dataset {
+  // Note that this Dataset API implementation uses an experimental tracing
+  // feature, which is not robust and does not have great diagnostics yet.
+  @inlinable
+  func map<ResultElement : TensorGroup>(
+    _ transform: (Element) -> ResultElement
+  ) -> Dataset<ResultElement> {
+    return Dataset<ResultElement>(_handle: Raw.mapDataset(
+      inputDataset: _handle,
+      otherArguments: Tensor<Int32>(0),
+      f: transform,
+      outputTypes: ResultElement._typeList,
+      outputShapes: ResultElement._unknownShapeList,
+      useInterOpParallelism: true,
+      preserveCardinality: false))
+  }
+
+  @inlinable
+  func map<ResultElement : TensorGroup>(
+    parallelCallCount: Int,
+    _ transform: (Element) -> ResultElement
+  ) -> Dataset<ResultElement> {
+    return Dataset<ResultElement>(_handle: Raw.parallelMapDataset(
+      inputDataset: _handle,
+      otherArguments: Tensor<Int32>(0),
+      numParallelCalls: Tensor<Int32>(Int32(parallelCallCount)),
+      f: transform,
+      outputTypes: ResultElement._typeList,
+      outputShapes: ResultElement._unknownShapeList,
+      useInterOpParallelism: true,
+      sloppy: false,
+      preserveCardinality: false))
+  }
+
+  @inlinable
+  func filter(
+    _ isIncluded: (Element) -> Tensor<Bool>
+  ) -> Dataset {
+    return Dataset(_handle: Raw.filterDataset(
+      inputDataset: _handle,
+      otherArguments: Tensor<Int32>(0),
+      predicate: isIncluded,
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList))
+  }
+}
+
+public extension Dataset {
+  @inlinable
+  func shuffled(
+    sampleCount: Int, randomSeed: Int64
+  ) -> Dataset {
+    let (seed1, seed2) = _tensorSeeds(Tensor(randomSeed))
+    return Dataset(_handle: Raw.shuffleDataset(
+      inputDataset: _handle,
+      bufferSize: Tensor(Int64(sampleCount)),
+      seed: seed1,
+      seed2: seed2,
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList))
+  }
+
+  @inlinable
+  func batched(_ batchSize: Int) -> Dataset {
+    return Dataset(_handle: Raw.batchDataset(
+      inputDataset: _handle,
+      batchSize: Tensor(Int64(batchSize)),
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList))
+  }
+}
+
+/// The type that allows iteration over a dataset's elements.
+@_fixed_layout
+public struct DatasetIterator<Element : TensorGroup> {
+  @usableFromInline let _handle: ResourceHandle
+
+  @usableFromInline
+  internal init(_handle: ResourceHandle) {
+    self._handle = _handle
+  }
+}
+
+extension DatasetIterator : IteratorProtocol {
+  /// Advances to the next element and returns it, or `nil` if no next element
+  /// exists.
+  @inlinable
+  public mutating func next() -> Element? {
+    let optional = Raw.iteratorGetNextAsOptional(
+      iterator: _handle,
+      outputTypes: Element._typeList,
+      outputShapes: Element._unknownShapeList)
+    guard Raw.optionalHasValue(optional: optional).scalarized() else {
+      return nil
+    }
+    return Raw.optionalGetValue(
+      optional: optional,
+      outputShapes: Element._unknownShapeList)
+  }
+}
+
+/// A 2-tuple-like struct that conforms to TensorGroup that represents a tuple
+/// of 2 types conforming to TensorGroup.
+@_fixed_layout
+public struct Zip2TensorGroup<T : TensorGroup, U : TensorGroup> : TensorGroup {
+  public var first: T
+  public var second: U
+
+  public init(_ first: T, _ second: U) {
+    self.first = first
+    self.second = second
+  }
+}
+
+@inlinable
+public func zip<T : TensorGroup, U : TensorGroup>(
+  _ dataset1: Dataset<T>, _ dataset2: Dataset<U>
+) -> Dataset<Zip2TensorGroup<T, U>> {
+  let handle = Raw.zipDataset(
+    inputDatasets: [dataset1._handle, dataset2._handle],
+    outputTypes: Zip2TensorGroup<T, U>._typeList,
+    outputShapes: Zip2TensorGroup<T, U>._unknownShapeList)
+  return Dataset(_handle: handle)
+}

Sources/DeepLearning/Layer.swift

@@ -513,6 +513,122 @@ public extension Conv2D {
     }
 }
 
+/// A 3-D convolution layer for spatial/spatio-temporal convolution over images.
+///
+/// This layer creates a convolution filter that is convolved with the layer input to produce a
+/// tensor of outputs.
+@_fixed_layout
+public struct Conv3D<Scalar: TensorFlowFloatingPoint>: Layer {
+    /// The 5-D convolution kernel.
+    public var filter: Tensor<Scalar>
+    /// The bias vector.
+    public var bias: Tensor<Scalar>
+    /// An activation function.
+    public typealias Activation = @differentiable (Tensor<Scalar>) -> Tensor<Scalar>
+    /// The element-wise activation function.
+    @noDerivative public let activation: Activation
+    /// The strides of the sliding window for spatial dimensions.
+    @noDerivative public let strides: (Int, Int, Int)
+    /// The padding algorithm for convolution.
+    @noDerivative public let padding: Padding
+
+    /// Creates a `Conv3D` layer with the specified filter, bias, activation function, strides, and
+    /// padding.
+    ///
+    /// - Parameters:
+    ///   - filter: The 5-D convolution kernel.
+    ///   - bias: The bias vector.
+    ///   - activation: The element-wise activation function.
+    ///   - strides: The strides of the sliding window for spatial dimensions.
+    ///   - padding: The padding algorithm for convolution.
+    public init(
+        filter: Tensor<Scalar>,
+        bias: Tensor<Scalar>,
+        activation: @escaping Activation,
+        strides: (Int, Int, Int),
+        padding: Padding
+    ) {
+        self.filter = filter
+        self.bias = bias
+        self.activation = activation
+        self.strides = strides
+        self.padding = padding
+    }
+
+    /// Returns the output obtained from applying the layer to the given input.
+    ///
+    /// - Parameter input: The input to the layer.
+    /// - Returns: The output.
+    @differentiable
+    public func call(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
+        return activation(input.convolved3D(withFilter: filter,
+                                            strides: (1, strides.0, strides.1, strides.2, 1),
+                                            padding: padding) + bias)
+    }
+}
+
+public extension Conv3D {
+    /// Creates a `Conv3D` layer with the specified filter shape, strides, padding, and
+    /// element-wise activation function. The filter tensor is initialized using Glorot uniform
+    /// initialization with the specified generator. The bias vector is initialized with zeros.
+    ///
+    /// - Parameters:
+    ///   - filterShape: The shape of the 5-D convolution kernel.
+    ///   - strides: The strides of the sliding window for spatial/spatio-temporal dimensions.
+    ///   - padding: The padding algorithm for convolution.
+    ///   - activation: The element-wise activation function.
+    ///   - generator: The random number generator for initialization.
+    ///
+    /// - Note: Use `init(filterShape:strides:padding:activation:seed:)` for faster random
+    ///   initialization.
+    init<G: RandomNumberGenerator>(
+        filterShape: (Int, Int, Int, Int, Int),
+        strides: (Int, Int, Int) = (1, 1, 1),
+        padding: Padding = .valid,
+        activation: @escaping Activation = identity,
+        generator: inout G
+    ) {
+        let filterTensorShape = TensorShape([
+            filterShape.0, filterShape.1, filterShape.2, filterShape.3, filterShape.4])
+        self.init(
+            filter: Tensor(glorotUniform: filterTensorShape, generator: &generator),
+            bias: Tensor(zeros: TensorShape([filterShape.4])),
+            activation: activation,
+            strides: strides,
+            padding: padding)
+    }
+}
+
+public extension Conv3D {
+    /// Creates a `Conv3D` layer with the specified filter shape, strides, padding, and
+    /// element-wise activation function. The filter tensor is initialized using Glorot uniform
+    /// initialization with the specified seed. The bias vector is initialized with zeros.
+    ///
+    /// - Parameters:
+    ///   - filterShape: The shape of the 5-D convolution kernel.
+    ///   - strides: The strides of the sliding window for spatial/spatio-temporal dimensions.
+    ///   - padding: The padding algorithm for convolution.
+    ///   - activation: The element-wise activation function.
+    ///   - seed: The random seed for initialization. The default value is random.
+    init(
+        filterShape: (Int, Int, Int, Int, Int),
+        strides: (Int, Int, Int) = (1, 1, 1),
+        padding: Padding = .valid,
+        activation: @escaping Activation = identity,
+        seed: (Int64, Int64) = (Int64.random(in: Int64.min..<Int64.max),
+                                Int64.random(in: Int64.min..<Int64.max))
+    ) {
+        let filterTensorShape = TensorShape([
+            filterShape.0, filterShape.1, filterShape.2, filterShape.3, filterShape.4])
+        self.init(
+            filter: Tensor(glorotUniform: filterTensorShape, seed: seed),
+            bias: Tensor(zeros: TensorShape([filterShape.4])),
+            activation: activation,
+            strides: strides,
+            padding: padding)
+    }
+}
+
 /// A 2-D transposed convolution layer (e.g. spatial transposed convolution over images).
 ///
 /// This layer creates a convolution filter that is transpose-convolved with the layer input
@@ -1224,7 +1340,7 @@ public struct UpSampling1D<Scalar: TensorFlowFloatingPoint>: Layer {
     ///
     /// - Parameter size: The upsampling factor for timesteps.
     public init(size: Int) {
-       self.size = size
+        self.size = size
     }
 
     /// Returns the output obtained from applying the layer to the given input.
@@ -1250,7 +1366,7 @@ public struct UpSampling2D<Scalar: TensorFlowFloatingPoint>: Layer {
     ///
     /// - Parameter size: The upsampling factor for rows and columns.
     public init(size: Int) {
-       self.size = size
+        self.size = size
     }
 
     /// Returns the output obtained from applying the layer to the given input.
@@ -1276,7 +1392,7 @@ public struct UpSampling3D<Scalar: TensorFlowFloatingPoint>: Layer {
     ///
     /// - Parameter size: The upsampling factor for rows and columns.
     public init(size: Int) {
-       self.size = size
+        self.size = size
     }
 
     /// Returns the output obtained from applying the layer to the given input.