Upsampling.swift

// Copyright 2019 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import _Differentiation

/// An upsampling layer for 1-D inputs.
@frozen
public struct UpSampling1D<Scalar: TensorFlowFloatingPoint>: ParameterlessLayer {
  public typealias TangentVector = EmptyTangentVector

  @noDerivative public let size: Int

  /// Creates an upsampling layer.
  ///
  /// - Parameter size: The upsampling factor for timesteps.
  public init(size: Int) {
    self.size = size
  }

  /// 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 callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
    let shape = input.shape
    let (batchSize, timesteps, channels) = (shape[0], shape[1], shape[2])
    let scaleOnes = Tensor<Scalar>(ones: [1, 1, size, 1], on: input.device)
    let upSampling = input.reshaped(to: [batchSize, timesteps, 1, channels]) * scaleOnes
    return upSampling.reshaped(to: [batchSize, timesteps * size, channels])
  }
}

/// An upsampling layer for 2-D inputs.
@frozen
public struct UpSampling2D<Scalar: TensorFlowFloatingPoint>: ParameterlessLayer {
  public typealias TangentVector = EmptyTangentVector

  @noDerivative public let size: Int

  /// Creates an upsampling layer.
  ///
  /// - Parameter size: The upsampling factor for rows and columns.
  public init(size: Int) {
    self.size = size
  }

  /// 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 callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
    let device = input.device
    let shape = input.shape
    let (batchSize, height, width, channels) = (shape[0], shape[1], shape[2], shape[3])
    let scaleOnes = Tensor<Scalar>(ones: [1, 1, size, 1, size, 1], on: device)
    let upSampling = input.reshaped(to: [batchSize, height, 1, width, 1, channels]) * scaleOnes
    return upSampling.reshaped(to: [batchSize, height * size, width * size, channels])
  }
}

/// An upsampling layer for 3-D inputs.
@frozen
public struct UpSampling3D<Scalar: TensorFlowFloatingPoint>: ParameterlessLayer {
  public typealias TangentVector = EmptyTangentVector

  @noDerivative public let size: Int

  /// Creates an upsampling layer.
  ///
  /// - Parameter size: The upsampling factor for rows and columns.
  public init(size: Int) {
    self.size = size
  }

  /// Repeats the elements of a tensor along an axis, like `np.repeat`.
  /// Function adapted from `def repeat_elements`:
  /// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/backend.py
  @differentiable
  private func repeatingElements(
    _ input: Tensor<Scalar>, alongAxis axis: Int, count: Int
  ) -> Tensor<Scalar> {
    let splits = _Raw.split(
      splitDim: Tensor<Int32>(Int32(axis), on: input.device),
      value: input,
      numSplit: Int64(input.shape[axis]))
    let repeated = splits.flatMap { x in Array(repeating: x, count: count) }
    return Tensor<Scalar>(concatenating: repeated, alongAxis: axis)
  }

  @derivative(of: repeatingElements)
  private func _vjpRepeatingElements(
    _ input: Tensor<Scalar>, alongAxis axis: Int, count: Int
  ) -> (value: Tensor<Scalar>, pullback: (Tensor<Scalar>) -> (TangentVector, Tensor<Scalar>)) {
    let value = repeatingElements(input, alongAxis: axis, count: count)
    return (
      value,
      { v in
        let splits = _Raw.split(
          splitDim: Tensor<Int32>(Int32(axis), on: v.device),
          value: v,
          numSplit: Int64(input.shape[axis]))
        let summed = splits.map { x in x.sum(alongAxes: axis) }
        let concatenated = Tensor<Scalar>(concatenating: summed, alongAxis: axis)
        return (.zero, concatenated)
      }
    )
  }

  /// 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 callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
    var result = repeatingElements(input, alongAxis: 1, count: size)
    result = repeatingElements(result, alongAxis: 2, count: size)
    result = repeatingElements(result, alongAxis: 3, count: size)
    return result
  }
}