From 106e303022d54ec852b55a4cf07e374a8f205fe5 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Wed, 22 May 2019 20:14:06 -0700 Subject: [PATCH 1/2] Fix UpSampling3D implementation. Broadcasting does not support 8-D tensors. Copy Keras implementation: simulating broadcasting via splitting, duplication, and concatenation. --- Sources/DeepLearning/Layer.swift | 35 ++++++++++++++++++++---- Tests/DeepLearningTests/LayerTests.swift | 7 ++--- 2 files changed, 32 insertions(+), 10 deletions(-) diff --git a/Sources/DeepLearning/Layer.swift b/Sources/DeepLearning/Layer.swift index f7560d282..2852a0f1a 100644 --- a/Sources/DeepLearning/Layer.swift +++ b/Sources/DeepLearning/Layer.swift @@ -1395,6 +1395,32 @@ public struct UpSampling3D: Layer { 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(vjp: _vjpRepeatingElements) + private func repeatingElements( + _ input: Tensor, alongAxis axis: Int, count: Int + ) -> Tensor { + let splits = Raw.split(splitDim: Tensor(Int32(axis)), + value: input, numSplit: Int64(input.shape[axis])) + let repeated = splits.flatMap { x in Array(repeating: x, count: count) } + return Tensor(concatenating: repeated, alongAxis: axis) + } + + private func _vjpRepeatingElements( + _ input: Tensor, alongAxis axis: Int, count: Int + ) -> (Tensor, (Tensor) -> (AllDifferentiableVariables, Tensor)) { + let value = repeatingElements(input, alongAxis: axis, count: count) + return (value, { v in + let splits = Raw.split(splitDim: Tensor(Int32(axis)), + value: v, numSplit: Int64(input.shape[axis])) + let summed = splits.map { x in x.sum(alongAxes: axis) } + let concatenated = Tensor(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. @@ -1404,11 +1430,10 @@ public struct UpSampling3D: Layer { let shape = input.shape let (batchSize, height, width, depth, channels) = (shape[0], shape[1], shape[2], shape[3], shape[4]) - let scaleOnes = Tensor(ones: [1, 1, size, 1, size, 1, size, 1]) - let upSampling = input.reshaped( - to: [batchSize, height, 1, width, 1, depth, 1, channels]) * scaleOnes - return upSampling.reshaped( - to: [batchSize, height * size, width * size, depth * size, channels]) + var result = repeatingElements(input, alongAxis: 1, count: height * size) + result = repeatingElements(result, alongAxis: 2, count: width * size) + result = repeatingElements(result, alongAxis: 3, count: depth * size) + return result } } diff --git a/Tests/DeepLearningTests/LayerTests.swift b/Tests/DeepLearningTests/LayerTests.swift index 9fd998852..887b70e7c 100644 --- a/Tests/DeepLearningTests/LayerTests.swift +++ b/Tests/DeepLearningTests/LayerTests.swift @@ -134,15 +134,12 @@ final class LayerTests: XCTestCase { let size = 6 let layer = UpSampling3D(size: size) let input = Tensor(shape: [1, 4, 3, 2, 1], scalars: (0..<24).map(Float.init)) - // TODO(TF-525): Fix `UpSampling3D.call`. - // Broadcasting does not support tensors with high rank: - // Broadcast between [1,4,1,3,1,2,1,1] and [1,1,6,1,6,1,6,1] is not supported yet. - /* + // TODO: Fix bug in `UpSampling3D.repeatingElements`. + // Result has shape `[1, 96, 54, 24, 1]` instead of `[1, 24, 18, 12, 1]`. let output = layer.inferring(from: input) let expected = TensorShape([1, input.shape[1] * size, input.shape[2] * size, input.shape[3] * size, 1]) XCTAssertEqual(output.shape, expected) XCTAssertEqual(output.shape, expected) - */ } func testReshape() { From c6e817c2c3b3d47d7535492fbc957ee4d21c4128 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Wed, 22 May 2019 20:51:46 -0700 Subject: [PATCH 2/2] Fix "repeating count" bug. --- Sources/DeepLearning/Layer.swift | 6 +++--- Tests/DeepLearningTests/LayerTests.swift | 2 -- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/Sources/DeepLearning/Layer.swift b/Sources/DeepLearning/Layer.swift index 2852a0f1a..e86887e8b 100644 --- a/Sources/DeepLearning/Layer.swift +++ b/Sources/DeepLearning/Layer.swift @@ -1430,9 +1430,9 @@ public struct UpSampling3D: Layer { let shape = input.shape let (batchSize, height, width, depth, channels) = (shape[0], shape[1], shape[2], shape[3], shape[4]) - var result = repeatingElements(input, alongAxis: 1, count: height * size) - result = repeatingElements(result, alongAxis: 2, count: width * size) - result = repeatingElements(result, alongAxis: 3, count: depth * size) + var result = repeatingElements(input, alongAxis: 1, count: size) + result = repeatingElements(result, alongAxis: 2, count: size) + result = repeatingElements(result, alongAxis: 3, count: size) return result } } diff --git a/Tests/DeepLearningTests/LayerTests.swift b/Tests/DeepLearningTests/LayerTests.swift index 887b70e7c..9fffb0866 100644 --- a/Tests/DeepLearningTests/LayerTests.swift +++ b/Tests/DeepLearningTests/LayerTests.swift @@ -134,8 +134,6 @@ final class LayerTests: XCTestCase { let size = 6 let layer = UpSampling3D(size: size) let input = Tensor(shape: [1, 4, 3, 2, 1], scalars: (0..<24).map(Float.init)) - // TODO: Fix bug in `UpSampling3D.repeatingElements`. - // Result has shape `[1, 96, 54, 24, 1]` instead of `[1, 24, 18, 12, 1]`. let output = layer.inferring(from: input) let expected = TensorShape([1, input.shape[1] * size, input.shape[2] * size, input.shape[3] * size, 1]) XCTAssertEqual(output.shape, expected)