From 17cd3c6ebceadb6c8f141c39e91a6c098763f4dd Mon Sep 17 00:00:00 2001 From: Ann Yuan Date: Wed, 13 May 2020 15:03:04 -0400 Subject: [PATCH 1/4] it begins --- tfjs-core/src/ops/array_ops.ts | 77 -------------------- tfjs-core/src/ops/batch_to_space_nd.ts | 99 ++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 77 deletions(-) create mode 100644 tfjs-core/src/ops/batch_to_space_nd.ts diff --git a/tfjs-core/src/ops/array_ops.ts b/tfjs-core/src/ops/array_ops.ts index 7164a8f4f00..bd915eaa248 100644 --- a/tfjs-core/src/ops/array_ops.ts +++ b/tfjs-core/src/ops/array_ops.ts @@ -159,83 +159,6 @@ function stack_( return concat(expandedTensors, axis); } -/** - * This operation reshapes the "batch" dimension 0 into `M + 1` dimensions of - * shape `blockShape + [batch]`, interleaves these blocks back into the grid - * defined by the spatial dimensions `[1, ..., M]`, to obtain a result with - * the same rank as the input. The spatial dimensions of this intermediate - * result are then optionally cropped according to `crops` to produce the - * output. This is the reverse of `tf.spaceToBatchND`. See below for a precise - * description. - * - * ```js - * const x = tf.tensor4d([1, 2, 3, 4], [4, 1, 1, 1]); - * const blockShape = [2, 2]; - * const crops = [[0, 0], [0, 0]]; - * - * x.batchToSpaceND(blockShape, crops).print(); - * ``` - * - * @param x A `tf.Tensor`. N-D with `x.shape` = `[batch] + spatialShape + - * remainingShape`, where spatialShape has `M` dimensions. - * @param blockShape A 1-D array. Must have shape `[M]`, all values must - * be >= 1. - * @param crops A 2-D array. Must have shape `[M, 2]`, all values must be >= 0. - * `crops[i] = [cropStart, cropEnd]` specifies the amount to crop from input - * dimension `i + 1`, which corresponds to spatial dimension `i`. It is required - * that `cropStart[i] + cropEnd[i] <= blockShape[i] * inputShape[i + 1]` - * - * This operation is equivalent to the following steps: - * - * 1. Reshape `x` to `reshaped` of shape: `[blockShape[0], ..., - * blockShape[M-1], batch / prod(blockShape), x.shape[1], ..., - * x.shape[N-1]]` - * - * 2. Permute dimensions of `reshaped`to produce `permuted` of shape `[batch / - * prod(blockShape),x.shape[1], blockShape[0], ..., x.shape[M], - * blockShape[M-1],x.shape[M+1], ..., x.shape[N-1]]` - * - * 3. Reshape `permuted` to produce `reshapedPermuted` of shape `[batch / - * prod(blockShape),x.shape[1] * blockShape[0], ..., x.shape[M] * - * blockShape[M-1],x.shape[M+1], ..., x.shape[N-1]]` - * - * 4. Crop the start and end of dimensions `[1, ..., M]` of `reshapedPermuted` - * according to `crops` to produce the output of shape: `[batch / - * prod(blockShape),x.shape[1] * blockShape[0] - crops[0,0] - crops[0,1], - * ..., x.shape[M] * blockShape[M-1] - crops[M-1,0] - - * crops[M-1,1],x.shape[M+1], ..., x.shape[N-1]]` - */ -/** @doc {heading: 'Tensors', subheading: 'Transformations'} */ -function batchToSpaceND_( - x: T|TensorLike, blockShape: number[], crops: number[][]): T { - const $x = convertToTensor(x, 'x', 'batchToSpaceND'); - const prod = blockShape.reduce((a, b) => a * b); - - util.assert( - $x.rank >= 1 + blockShape.length, - () => `input rank is ${$x.rank} but should be > than blockShape.length ${ - blockShape.length}`); - - util.assert( - crops.length === blockShape.length, - () => `crops.length is ${ - crops.length} but should be equal to blockShape.length ${ - blockShape.length}`); - - util.assert( - $x.shape[0] % prod === 0, - () => `input tensor batch is ${ - $x.shape[0]} but is not divisible by the product of ` + - `the elements of blockShape ${blockShape.join(' * ')} === ${prod}`); - - const grad = (dy: T) => { - return {$x: () => dy.spaceToBatchND(blockShape, crops)}; - }; - - return ENGINE.runKernelFunc( - backend => backend.batchToSpaceND($x, blockShape, crops), {$x}, grad); -} - /** * This operation divides "spatial" dimensions `[1, ..., M]` of the input into * a grid of blocks of shape `blockShape`, and interleaves these blocks with diff --git a/tfjs-core/src/ops/batch_to_space_nd.ts b/tfjs-core/src/ops/batch_to_space_nd.ts new file mode 100644 index 00000000000..0d49a3cae2a --- /dev/null +++ b/tfjs-core/src/ops/batch_to_space_nd.ts @@ -0,0 +1,99 @@ +/** + * @license + * Copyright 2020 Google Inc. 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 {ENGINE} from '../engine'; +import {Tensor, Tensor4D, TensorBuffer} from '../tensor'; +import {convertToTensor, convertToTensorArray} from '../tensor_util_env'; +import {TensorLike} from '../types'; +import * as util from '../util'; + +/** + * This operation reshapes the "batch" dimension 0 into `M + 1` dimensions of + * shape `blockShape + [batch]`, interleaves these blocks back into the grid + * defined by the spatial dimensions `[1, ..., M]`, to obtain a result with + * the same rank as the input. The spatial dimensions of this intermediate + * result are then optionally cropped according to `crops` to produce the + * output. This is the reverse of `tf.spaceToBatchND`. See below for a precise + * description. + * + * ```js + * const x = tf.tensor4d([1, 2, 3, 4], [4, 1, 1, 1]); + * const blockShape = [2, 2]; + * const crops = [[0, 0], [0, 0]]; + * + * x.batchToSpaceND(blockShape, crops).print(); + * ``` + * + * @param x A `tf.Tensor`. N-D with `x.shape` = `[batch] + spatialShape + + * remainingShape`, where spatialShape has `M` dimensions. + * @param blockShape A 1-D array. Must have shape `[M]`, all values must + * be >= 1. + * @param crops A 2-D array. Must have shape `[M, 2]`, all values must be >= 0. + * `crops[i] = [cropStart, cropEnd]` specifies the amount to crop from input + * dimension `i + 1`, which corresponds to spatial dimension `i`. It is required + * that `cropStart[i] + cropEnd[i] <= blockShape[i] * inputShape[i + 1]` + * + * This operation is equivalent to the following steps: + * + * 1. Reshape `x` to `reshaped` of shape: `[blockShape[0], ..., + * blockShape[M-1], batch / prod(blockShape), x.shape[1], ..., + * x.shape[N-1]]` + * + * 2. Permute dimensions of `reshaped`to produce `permuted` of shape `[batch / + * prod(blockShape),x.shape[1], blockShape[0], ..., x.shape[M], + * blockShape[M-1],x.shape[M+1], ..., x.shape[N-1]]` + * + * 3. Reshape `permuted` to produce `reshapedPermuted` of shape `[batch / + * prod(blockShape),x.shape[1] * blockShape[0], ..., x.shape[M] * + * blockShape[M-1],x.shape[M+1], ..., x.shape[N-1]]` + * + * 4. Crop the start and end of dimensions `[1, ..., M]` of `reshapedPermuted` + * according to `crops` to produce the output of shape: `[batch / + * prod(blockShape),x.shape[1] * blockShape[0] - crops[0,0] - crops[0,1], + * ..., x.shape[M] * blockShape[M-1] - crops[M-1,0] - + * crops[M-1,1],x.shape[M+1], ..., x.shape[N-1]]` + */ +/** @doc {heading: 'Tensors', subheading: 'Transformations'} */ +function batchToSpaceND_( + x: T|TensorLike, blockShape: number[], crops: number[][]): T { + const $x = convertToTensor(x, 'x', 'batchToSpaceND'); + const prod = blockShape.reduce((a, b) => a * b); + + util.assert( + $x.rank >= 1 + blockShape.length, + () => `input rank is ${$x.rank} but should be > than blockShape.length ${ + blockShape.length}`); + + util.assert( + crops.length === blockShape.length, + () => `crops.length is ${ + crops.length} but should be equal to blockShape.length ${ + blockShape.length}`); + + util.assert( + $x.shape[0] % prod === 0, + () => `input tensor batch is ${ + $x.shape[0]} but is not divisible by the product of ` + + `the elements of blockShape ${blockShape.join(' * ')} === ${prod}`); + + const grad = (dy: T) => { + return {$x: () => dy.spaceToBatchND(blockShape, crops)}; + }; + + return ENGINE.runKernelFunc( + backend => backend.batchToSpaceND($x, blockShape, crops), {$x}, grad); +} From 0f874d61f7ec57b861f4136275665cf6da7765a3 Mon Sep 17 00:00:00 2001 From: Ann Yuan Date: Thu, 14 May 2020 11:37:07 -0400 Subject: [PATCH 2/4] modularise --- .../src/gradients/BatchToSpaceND_grad.ts | 31 +++ tfjs-core/src/kernel_names.ts | 7 + tfjs-core/src/ops/array_ops.ts | 1 - tfjs-core/src/ops/array_ops_test.ts | 194 ---------------- tfjs-core/src/ops/batch_to_space_nd.ts | 23 +- tfjs-core/src/ops/batch_to_space_nd_test.ts | 214 ++++++++++++++++++ tfjs-core/src/ops/ops.ts | 1 + tfjs-core/src/ops/pool.ts | 7 +- .../public/chained_ops/batch_to_space_nd.ts | 32 +++ .../chained_ops/register_all_chained_ops.ts | 1 + .../register_all_chained_ops_test.ts | 1 + tfjs-core/src/register_all_gradients.ts | 2 + tfjs-core/src/tensor.ts | 8 - 13 files changed, 310 insertions(+), 212 deletions(-) create mode 100644 tfjs-core/src/gradients/BatchToSpaceND_grad.ts create mode 100644 tfjs-core/src/ops/batch_to_space_nd_test.ts create mode 100644 tfjs-core/src/public/chained_ops/batch_to_space_nd.ts diff --git a/tfjs-core/src/gradients/BatchToSpaceND_grad.ts b/tfjs-core/src/gradients/BatchToSpaceND_grad.ts new file mode 100644 index 00000000000..ccf57c96849 --- /dev/null +++ b/tfjs-core/src/gradients/BatchToSpaceND_grad.ts @@ -0,0 +1,31 @@ +/** + * @license + * Copyright 2020 Google Inc. 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 {BatchToSpaceND, BatchToSpaceNDAttrs} from '../kernel_names'; +import {GradConfig, NamedAttrMap} from '../kernel_registry'; +import {spaceToBatchND} from '../ops/array_ops'; +import {Tensor} from '../tensor'; + +export const batchToSpaceNDGradConfig: GradConfig = { + kernelName: BatchToSpaceND, + gradFunc: (dy: Tensor, saved: Tensor[], attrs: NamedAttrMap) => { + const batchToSpaceNDAttrs: BatchToSpaceNDAttrs = + attrs as {} as BatchToSpaceNDAttrs; + const {blockShape, crops} = batchToSpaceNDAttrs; + return {x: () => spaceToBatchND(dy, blockShape, crops)}; + } +}; diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts index 9347cc68e89..06b61ec7915 100644 --- a/tfjs-core/src/kernel_names.ts +++ b/tfjs-core/src/kernel_names.ts @@ -34,6 +34,13 @@ export interface BatchMatMulAttrs { transposeB: boolean; } +export const BatchToSpaceND = 'BatchToSpaceND'; +export type BatchToSpaceNDInputs = Pick; +export interface BatchToSpaceNDAttrs { + blockShape: number[]; + crops: number[][]; +} + export type BinaryInputs = Pick; export const BroadcastTo = 'BroadcastTo'; diff --git a/tfjs-core/src/ops/array_ops.ts b/tfjs-core/src/ops/array_ops.ts index bd915eaa248..555fad9c871 100644 --- a/tfjs-core/src/ops/array_ops.ts +++ b/tfjs-core/src/ops/array_ops.ts @@ -543,7 +543,6 @@ export { print // Not wrapped in op() since no need to increase stack trace. }; -export const batchToSpaceND = op({batchToSpaceND_}); export const cast = op({cast_}); export const cumsum = op({cumsum_}); export const depthToSpace = op({depthToSpace_}); diff --git a/tfjs-core/src/ops/array_ops_test.ts b/tfjs-core/src/ops/array_ops_test.ts index 0311fbac485..60ed116b3d5 100644 --- a/tfjs-core/src/ops/array_ops_test.ts +++ b/tfjs-core/src/ops/array_ops_test.ts @@ -2852,200 +2852,6 @@ describeWithFlags('cumsum', ALL_ENVS, () => { }); }); -describeWithFlags('batchToSpaceND', ALL_ENVS, () => { - it('tensor4d, input shape=[4, 1, 1, 1], blockShape=[2, 2]', async () => { - const t = tf.tensor4d([1, 2, 3, 4], [4, 1, 1, 1]); - const blockShape = [2, 2]; - const crops = [[0, 0], [0, 0]]; - - const res = tf.batchToSpaceND(t, blockShape, crops); - expect(res.shape).toEqual([1, 2, 2, 1]); - expectArraysClose(await res.data(), [1, 2, 3, 4]); - }); - - it('tensor4d, input shape=[4, 1, 1, 3], blockShape=[2, 2]', async () => { - const t = - tf.tensor4d([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], [4, 1, 1, 3]); - const blockShape = [2, 2]; - const crops = [[0, 0], [0, 0]]; - - const res = tf.batchToSpaceND(t, blockShape, crops); - expect(res.shape).toEqual([1, 2, 2, 3]); - expectArraysClose( - await res.data(), [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]); - }); - - it('tensor4d, input shape=[4, 2, 2, 1], blockShape=[2, 2]', async () => { - const t = tf.tensor4d( - [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16], [4, 2, 2, 1]); - const blockShape = [2, 2]; - const crops = [[0, 0], [0, 0]]; - - const res = tf.batchToSpaceND(t, blockShape, crops); - expect(res.shape).toEqual([1, 4, 4, 1]); - expectArraysClose( - await res.data(), - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]); - }); - - it('tensor4d, input shape=[8, 1, 3, 1], blockShape=[2, 2]', async () => { - const t = tf.tensor4d( - [ - 0, 1, 3, 0, 9, 11, 0, 2, 4, 0, 10, 12, - 0, 5, 7, 0, 13, 15, 0, 6, 8, 0, 14, 16 - ], - [8, 1, 3, 1]); - const blockShape = [2, 2]; - const crops = [[0, 0], [2, 0]]; - - const res = tf.batchToSpaceND(t, blockShape, crops); - expect(res.shape).toEqual([2, 2, 4, 1]); - expectArraysClose( - await res.data(), - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]); - }); - - it('tensor2d, blockShape [1]', async () => { - const t = tf.tensor2d([1, 2, 3, 4], [2, 2]); - const blockShape = [2]; - const crops = [[0, 0]]; - - const res = tf.batchToSpaceND(t, blockShape, crops); - expect(res.shape).toEqual([1, 4]); - expectArraysClose(await res.data(), [1, 3, 2, 4]); - }); - - it('tensor3d, blockSHape [1]', async () => { - const t = tf.tensor( - [ - -61, 37, -68, 72, 31, 62, 0, -13, 28, 54, 96, - 44, -55, -64, -88, -94, 65, -32, -96, -73, -2, -77, - -14, 47, 33, 15, 70, 20, 75, 28, 84, -13 - ], - [8, 2, 2]); - const blockShape = [2]; - const crops = [[0, 2]]; - - const res = tf.batchToSpaceND(t, blockShape, crops); - expect(res.shape).toEqual([4, 2, 2]); - expectArraysClose( - await res.data(), - [-61, 37, 65, -32, 31, 62, -2, -77, 28, 54, 33, 15, -55, -64, 75, 28]); - }); - - it('tensor3d, blockShape [2]', async () => { - const t = tf.tensor( - [ - -61, 37, -68, 72, 31, 62, 0, -13, 28, 54, 96, - 44, -55, -64, -88, -94, 65, -32, -96, -73, -2, -77, - -14, 47, 33, 15, 70, 20, 75, 28, 84, -13 - ], - [8, 2, 2]); - const blockShape = [2, 2]; - const crops = [[2, 0], [2, 0]]; - - const res = tf.batchToSpaceND(t, blockShape, crops); - expect(res.shape).toEqual([2, 2, 2]); - expectArraysClose(await res.data(), [72, 44, -73, 20, -13, -94, 47, -13]); - }); - - it('throws when blockShape equal to input rank', () => { - const t = tf.tensor4d([1, 2, 3, 4], [4, 1, 1, 1]); - const blockShape = [2, 2, 2, 2]; - const crops = [[0, 0], [0, 0], [0, 0], [0, 0]]; - - expect(() => tf.batchToSpaceND(t, blockShape, crops)) - .toThrowError( - `input rank is ${t.rank} but should be > than blockShape.length ${ - blockShape.length}`); - }); - - it('throws when crops row dimension not equal to blockshape', () => { - const t = tf.tensor4d([1, 2, 3, 4], [4, 1, 1, 1]); - const blockShape = [2, 2]; - const crops = [[0, 0]]; - - expect(() => tf.batchToSpaceND(t, blockShape, crops)) - .toThrowError(`crops.length is ${ - crops.length} but should be equal to blockShape.length ${ - blockShape.length}`); - }); - - it('throws when input tensor batch not divisible by prod(blockShape)', () => { - const t = tf.tensor4d([1, 2, 3, 4, 5], [5, 1, 1, 1]); - const blockShape = [2, 2]; - const crops = [[0, 0], [0, 0]]; - const prod = blockShape.reduce((a, b) => a * b); - - expect(() => tf.batchToSpaceND(t, blockShape, crops)) - .toThrowError( - `input tensor batch is ${t.shape[0]} but is not divisible by the ` + - `product of the elements of blockShape ${ - blockShape.join(' * ')} === ${prod}`); - }); - - it('accepts a tensor-like object', async () => { - const t = [[[[1]]], [[[2]]], [[[3]]], [[[4]]]]; - const blockShape = [2, 2]; - const crops = [[0, 0], [0, 0]]; - - const res = tf.batchToSpaceND(t, blockShape, crops); - expect(res.shape).toEqual([1, 2, 2, 1]); - expectArraysClose(await res.data(), [1, 2, 3, 4]); - }); - - it('gradients, input shape=[4, 2, 2], block shape=[2]', async () => { - const t = tf.tensor( - [-61, 37, -68, 72, 31, 62, 0, -13, 28, 54, 96, 44, -55, -64, -88, -94], - [4, 2, 2]); - const blockShape = [2]; - const crops = [[0, 2]]; - const dy = tf.tensor([.01, .02, .03, .04, .05, .06, .07, .08], [2, 2, 2]); - - const gradient = - tf.grad(t => tf.batchToSpaceND(t, blockShape, crops))(t, dy); - expect(gradient.shape).toEqual([4, 2, 2]); - expectArraysClose(await gradient.data(), [ - 0.01, 0.02, 0, 0, 0.05, 0.06, 0, 0, 0.03, 0.04, 0, 0, 0.07, 0.08, 0, 0 - ]); - }); - - it('gradients, input shape=[4, 2, 2, 1], block shape=[2, 2]', async () => { - const t = tf.tensor4d( - [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16], [4, 2, 2, 1]); - const blockShape = [2, 2]; - const crops = [[0, 0], [0, 0]]; - const dy = tf.tensor( - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], [1, 4, 4, 1]); - - const gradient = - tf.grad(t => tf.batchToSpaceND(t, blockShape, crops))(t, dy); - expect(gradient.shape).toEqual([4, 2, 2, 1]); - expectArraysClose( - await gradient.data(), - [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16]); - }); - - it('gradient with clones, input=[4, 2, 2, 1], block shape=[2, 2]', - async () => { - const t = tf.tensor4d( - [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16], - [4, 2, 2, 1]); - const blockShape = [2, 2]; - const crops = [[0, 0], [0, 0]]; - const dy = tf.tensor( - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], - [1, 4, 4, 1]); - - const gradient = tf.grad( - t => tf.batchToSpaceND(t.clone(), blockShape, crops).clone())(t, dy); - expect(gradient.shape).toEqual([4, 2, 2, 1]); - expectArraysClose( - await gradient.data(), - [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16]); - }); -}); - describeWithFlags('spaceToBatchND', ALL_ENVS, () => { it('tensor4d, input shape=[1, 2, 2, 1], blockShape=[2, 2]', async () => { const t = tf.tensor4d([[[[1], [2]], [[3], [4]]]], [1, 2, 2, 1]); diff --git a/tfjs-core/src/ops/batch_to_space_nd.ts b/tfjs-core/src/ops/batch_to_space_nd.ts index 0d49a3cae2a..90cb0bd332b 100644 --- a/tfjs-core/src/ops/batch_to_space_nd.ts +++ b/tfjs-core/src/ops/batch_to_space_nd.ts @@ -15,12 +15,17 @@ * ============================================================================= */ -import {ENGINE} from '../engine'; -import {Tensor, Tensor4D, TensorBuffer} from '../tensor'; -import {convertToTensor, convertToTensorArray} from '../tensor_util_env'; +import {ENGINE, ForwardFunc} from '../engine'; +import {BatchToSpaceND, BatchToSpaceNDAttrs, BatchToSpaceNDInputs} from '../kernel_names'; +import {NamedAttrMap} from '../kernel_registry'; +import {Tensor} from '../tensor'; +import {NamedTensorMap} from '../tensor_types'; +import {convertToTensor} from '../tensor_util_env'; import {TensorLike} from '../types'; import * as util from '../util'; +import {op} from './operation'; + /** * This operation reshapes the "batch" dimension 0 into `M + 1` dimensions of * shape `blockShape + [batch]`, interleaves these blocks back into the grid @@ -90,10 +95,16 @@ function batchToSpaceND_( $x.shape[0]} but is not divisible by the product of ` + `the elements of blockShape ${blockShape.join(' * ')} === ${prod}`); - const grad = (dy: T) => { - return {$x: () => dy.spaceToBatchND(blockShape, crops)}; + const forward: ForwardFunc = backend => { + return backend.batchToSpaceND($x, blockShape, crops); }; + const inputs: BatchToSpaceNDInputs = {x: $x}; + const attrs: BatchToSpaceNDAttrs = {blockShape, crops}; + return ENGINE.runKernelFunc( - backend => backend.batchToSpaceND($x, blockShape, crops), {$x}, grad); + forward, inputs as {} as NamedTensorMap, null /* gradient */, + BatchToSpaceND, attrs as {} as NamedAttrMap); } + +export const batchToSpaceND = op({batchToSpaceND_}); diff --git a/tfjs-core/src/ops/batch_to_space_nd_test.ts b/tfjs-core/src/ops/batch_to_space_nd_test.ts new file mode 100644 index 00000000000..e003cebd8de --- /dev/null +++ b/tfjs-core/src/ops/batch_to_space_nd_test.ts @@ -0,0 +1,214 @@ +/** + * @license + * Copyright 2018 Google Inc. 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 * as tf from '../index'; +import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; +import {expectArraysClose} from '../test_util'; + +describeWithFlags('batchToSpaceND', ALL_ENVS, () => { + it('tensor4d, input shape=[4, 1, 1, 1], blockShape=[2, 2]', async () => { + const t = tf.tensor4d([1, 2, 3, 4], [4, 1, 1, 1]); + const blockShape = [2, 2]; + const crops = [[0, 0], [0, 0]]; + + const res = tf.batchToSpaceND(t, blockShape, crops); + expect(res.shape).toEqual([1, 2, 2, 1]); + expectArraysClose(await res.data(), [1, 2, 3, 4]); + }); + + it('tensor4d, input shape=[4, 1, 1, 3], blockShape=[2, 2]', async () => { + const t = + tf.tensor4d([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], [4, 1, 1, 3]); + const blockShape = [2, 2]; + const crops = [[0, 0], [0, 0]]; + + const res = tf.batchToSpaceND(t, blockShape, crops); + expect(res.shape).toEqual([1, 2, 2, 3]); + expectArraysClose( + await res.data(), [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]); + }); + + it('tensor4d, input shape=[4, 2, 2, 1], blockShape=[2, 2]', async () => { + const t = tf.tensor4d( + [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16], [4, 2, 2, 1]); + const blockShape = [2, 2]; + const crops = [[0, 0], [0, 0]]; + + const res = tf.batchToSpaceND(t, blockShape, crops); + expect(res.shape).toEqual([1, 4, 4, 1]); + expectArraysClose( + await res.data(), + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]); + }); + + it('tensor4d, input shape=[8, 1, 3, 1], blockShape=[2, 2]', async () => { + const t = tf.tensor4d( + [ + 0, 1, 3, 0, 9, 11, 0, 2, 4, 0, 10, 12, + 0, 5, 7, 0, 13, 15, 0, 6, 8, 0, 14, 16 + ], + [8, 1, 3, 1]); + const blockShape = [2, 2]; + const crops = [[0, 0], [2, 0]]; + + const res = tf.batchToSpaceND(t, blockShape, crops); + expect(res.shape).toEqual([2, 2, 4, 1]); + expectArraysClose( + await res.data(), + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]); + }); + + it('tensor2d, blockShape [1]', async () => { + const t = tf.tensor2d([1, 2, 3, 4], [2, 2]); + const blockShape = [2]; + const crops = [[0, 0]]; + + const res = tf.batchToSpaceND(t, blockShape, crops); + expect(res.shape).toEqual([1, 4]); + expectArraysClose(await res.data(), [1, 3, 2, 4]); + }); + + it('tensor3d, blockSHape [1]', async () => { + const t = tf.tensor( + [ + -61, 37, -68, 72, 31, 62, 0, -13, 28, 54, 96, + 44, -55, -64, -88, -94, 65, -32, -96, -73, -2, -77, + -14, 47, 33, 15, 70, 20, 75, 28, 84, -13 + ], + [8, 2, 2]); + const blockShape = [2]; + const crops = [[0, 2]]; + + const res = tf.batchToSpaceND(t, blockShape, crops); + expect(res.shape).toEqual([4, 2, 2]); + expectArraysClose( + await res.data(), + [-61, 37, 65, -32, 31, 62, -2, -77, 28, 54, 33, 15, -55, -64, 75, 28]); + }); + + it('tensor3d, blockShape [2]', async () => { + const t = tf.tensor( + [ + -61, 37, -68, 72, 31, 62, 0, -13, 28, 54, 96, + 44, -55, -64, -88, -94, 65, -32, -96, -73, -2, -77, + -14, 47, 33, 15, 70, 20, 75, 28, 84, -13 + ], + [8, 2, 2]); + const blockShape = [2, 2]; + const crops = [[2, 0], [2, 0]]; + + const res = tf.batchToSpaceND(t, blockShape, crops); + expect(res.shape).toEqual([2, 2, 2]); + expectArraysClose(await res.data(), [72, 44, -73, 20, -13, -94, 47, -13]); + }); + + it('throws when blockShape equal to input rank', () => { + const t = tf.tensor4d([1, 2, 3, 4], [4, 1, 1, 1]); + const blockShape = [2, 2, 2, 2]; + const crops = [[0, 0], [0, 0], [0, 0], [0, 0]]; + + expect(() => tf.batchToSpaceND(t, blockShape, crops)) + .toThrowError( + `input rank is ${t.rank} but should be > than blockShape.length ${ + blockShape.length}`); + }); + + it('throws when crops row dimension not equal to blockshape', () => { + const t = tf.tensor4d([1, 2, 3, 4], [4, 1, 1, 1]); + const blockShape = [2, 2]; + const crops = [[0, 0]]; + + expect(() => tf.batchToSpaceND(t, blockShape, crops)) + .toThrowError(`crops.length is ${ + crops.length} but should be equal to blockShape.length ${ + blockShape.length}`); + }); + + it('throws when input tensor batch not divisible by prod(blockShape)', () => { + const t = tf.tensor4d([1, 2, 3, 4, 5], [5, 1, 1, 1]); + const blockShape = [2, 2]; + const crops = [[0, 0], [0, 0]]; + const prod = blockShape.reduce((a, b) => a * b); + + expect(() => tf.batchToSpaceND(t, blockShape, crops)) + .toThrowError( + `input tensor batch is ${t.shape[0]} but is not divisible by the ` + + `product of the elements of blockShape ${ + blockShape.join(' * ')} === ${prod}`); + }); + + it('accepts a tensor-like object', async () => { + const t = [[[[1]]], [[[2]]], [[[3]]], [[[4]]]]; + const blockShape = [2, 2]; + const crops = [[0, 0], [0, 0]]; + + const res = tf.batchToSpaceND(t, blockShape, crops); + expect(res.shape).toEqual([1, 2, 2, 1]); + expectArraysClose(await res.data(), [1, 2, 3, 4]); + }); + + it('gradients, input shape=[4, 2, 2], block shape=[2]', async () => { + const t = tf.tensor( + [-61, 37, -68, 72, 31, 62, 0, -13, 28, 54, 96, 44, -55, -64, -88, -94], + [4, 2, 2]); + const blockShape = [2]; + const crops = [[0, 2]]; + const dy = tf.tensor([.01, .02, .03, .04, .05, .06, .07, .08], [2, 2, 2]); + + const gradient = + tf.grad(t => tf.batchToSpaceND(t, blockShape, crops))(t, dy); + expect(gradient.shape).toEqual([4, 2, 2]); + expectArraysClose(await gradient.data(), [ + 0.01, 0.02, 0, 0, 0.05, 0.06, 0, 0, 0.03, 0.04, 0, 0, 0.07, 0.08, 0, 0 + ]); + }); + + it('gradients, input shape=[4, 2, 2, 1], block shape=[2, 2]', async () => { + const t = tf.tensor4d( + [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16], [4, 2, 2, 1]); + const blockShape = [2, 2]; + const crops = [[0, 0], [0, 0]]; + const dy = tf.tensor( + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], [1, 4, 4, 1]); + + const gradient = + tf.grad(t => tf.batchToSpaceND(t, blockShape, crops))(t, dy); + expect(gradient.shape).toEqual([4, 2, 2, 1]); + expectArraysClose( + await gradient.data(), + [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16]); + }); + + it('gradient with clones, input=[4, 2, 2, 1], block shape=[2, 2]', + async () => { + const t = tf.tensor4d( + [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16], + [4, 2, 2, 1]); + const blockShape = [2, 2]; + const crops = [[0, 0], [0, 0]]; + const dy = tf.tensor( + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + [1, 4, 4, 1]); + + const gradient = tf.grad( + t => tf.batchToSpaceND(t.clone(), blockShape, crops).clone())(t, dy); + expect(gradient.shape).toEqual([4, 2, 2, 1]); + expectArraysClose( + await gradient.data(), + [1, 3, 9, 11, 2, 4, 10, 12, 5, 7, 13, 15, 6, 8, 14, 16]); + }); +}); diff --git a/tfjs-core/src/ops/ops.ts b/tfjs-core/src/ops/ops.ts index 3fe8696e938..9cc1e31d9f2 100644 --- a/tfjs-core/src/ops/ops.ts +++ b/tfjs-core/src/ops/ops.ts @@ -18,6 +18,7 @@ // Modularized ops. export {add} from './add'; export {addN} from './add_n'; +export {batchToSpaceND} from './batch_to_space_nd'; export {batchNorm} from './batchnorm'; export {batchNorm2d} from './batchnorm2d'; export {batchNorm3d} from './batchnorm3d'; diff --git a/tfjs-core/src/ops/pool.ts b/tfjs-core/src/ops/pool.ts index 8bfd8fc33b0..6063385838f 100644 --- a/tfjs-core/src/ops/pool.ts +++ b/tfjs-core/src/ops/pool.ts @@ -22,7 +22,8 @@ import {convertToTensor} from '../tensor_util_env'; import {TensorLike} from '../types'; import * as util from '../util'; -import {batchToSpaceND, spaceToBatchND} from './array_ops'; +import {spaceToBatchND} from './array_ops'; +import {batchToSpaceND} from './batch_to_space_nd'; import * as conv_util from './conv_util'; import {op} from './operation'; @@ -917,8 +918,8 @@ function maxPool3dBackprop( /** @doc {heading: 'Operations', subheading: 'Convolution'} */ function maxPoolWithArgmax_( x: T|TensorLike, filterSize: [number, number]|number, - strides: [number, number]|number, - pad: 'valid'|'same'|number, includeBatchInIndex = false): NamedTensorMap { + strides: [number, number]|number, pad: 'valid'|'same'|number, + includeBatchInIndex = false): NamedTensorMap { const $x = convertToTensor(x, 'x', 'maxPoolWithArgmax'); const attrs = {filterSize, strides, pad, includeBatchInIndex}; diff --git a/tfjs-core/src/public/chained_ops/batch_to_space_nd.ts b/tfjs-core/src/public/chained_ops/batch_to_space_nd.ts new file mode 100644 index 00000000000..0a4160a5968 --- /dev/null +++ b/tfjs-core/src/public/chained_ops/batch_to_space_nd.ts @@ -0,0 +1,32 @@ +/** + * @license + * Copyright 2020 Google LLC. 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 {batchToSpaceND} from '../../ops/batch_to_space_nd'; +import {Tensor} from '../../tensor'; +import {Rank} from '../../types'; + +declare module '../../tensor' { + interface Tensor { + batchToSpaceND(blockShape: number[], crops: number[][]): + Tensor; + } +} + +Tensor.prototype.batchToSpaceND = function( + blockShape: number[], crops: number[][]): Tensor { + this.throwIfDisposed(); + return batchToSpaceND(this, blockShape, crops); +}; diff --git a/tfjs-core/src/public/chained_ops/register_all_chained_ops.ts b/tfjs-core/src/public/chained_ops/register_all_chained_ops.ts index e2134bd8204..e8a285a4954 100644 --- a/tfjs-core/src/public/chained_ops/register_all_chained_ops.ts +++ b/tfjs-core/src/public/chained_ops/register_all_chained_ops.ts @@ -17,6 +17,7 @@ import './add'; import './batchnorm'; import './broadcast_to'; +import './batch_to_space_nd'; import './max'; import './concat'; import './conv1d'; diff --git a/tfjs-core/src/public/chained_ops/register_all_chained_ops_test.ts b/tfjs-core/src/public/chained_ops/register_all_chained_ops_test.ts index 776da604145..67cb0ce55e6 100644 --- a/tfjs-core/src/public/chained_ops/register_all_chained_ops_test.ts +++ b/tfjs-core/src/public/chained_ops/register_all_chained_ops_test.ts @@ -26,6 +26,7 @@ import {ALL_ENVS, describeWithFlags} from '../../jasmine_util'; const CHAINED_OPS = [ 'add', 'batchNorm', + 'batchToSpaceND', 'broadcastTo', 'concat', 'conv1d', diff --git a/tfjs-core/src/register_all_gradients.ts b/tfjs-core/src/register_all_gradients.ts index 2526cf2211c..00248c5e305 100644 --- a/tfjs-core/src/register_all_gradients.ts +++ b/tfjs-core/src/register_all_gradients.ts @@ -17,6 +17,7 @@ import {addGradConfig} from './gradients/Add_grad'; import {addNGradConfig} from './gradients/AddN_grad'; import {batchMatMulGradConfig} from './gradients/BatchMatMul_grad'; +import {batchToSpaceNDGradConfig} from './gradients/BatchToSpaceND_grad'; import {broadcastToGradConfig} from './gradients/BroadcastTo_grad'; import {concatGradConfig} from './gradients/Concat_grad'; import {conv2DGradConfig} from './gradients/Conv2D_grad'; @@ -45,6 +46,7 @@ const gradConfigs: GradConfig[] = [ addGradConfig, addNGradConfig, batchMatMulGradConfig, + batchToSpaceNDGradConfig, broadcastToGradConfig, concatGradConfig, conv2DGradConfig, diff --git a/tfjs-core/src/tensor.ts b/tfjs-core/src/tensor.ts index f83955cbd24..446c83b3c01 100644 --- a/tfjs-core/src/tensor.ts +++ b/tfjs-core/src/tensor.ts @@ -295,8 +295,6 @@ export interface OpHandler { strides?: [number, number]|number): T; unsortedSegmentSum( x: T, segmentIds: Tensor1D|TensorLike1D, numSegments: number): T; - batchToSpaceND( - x: T, blockShape: number[], crops: number[][]): T; spaceToBatchND( x: T, blockShape: number[], paddings: number[][]): T; topk(x: T, k: number, sorted: boolean): @@ -1184,12 +1182,6 @@ export class Tensor { return opHandler.unsortedSegmentSum(this, segmentIds, numSegments); } - batchToSpaceND( - this: T, blockShape: number[], crops: number[][]): T { - this.throwIfDisposed(); - return opHandler.batchToSpaceND(this, blockShape, crops); - } - spaceToBatchND( this: T, blockShape: number[], paddings: number[][]): T { this.throwIfDisposed(); From a5e0c0d61942a018ced14fe878a3ef00e1c0e039 Mon Sep 17 00:00:00 2001 From: Ann Yuan Date: Thu, 14 May 2020 11:58:33 -0400 Subject: [PATCH 3/4] types --- tfjs-core/src/ops/array_ops.ts | 2 +- tfjs-core/src/tests.ts | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tfjs-core/src/ops/array_ops.ts b/tfjs-core/src/ops/array_ops.ts index 555fad9c871..5c6ec4a2a0d 100644 --- a/tfjs-core/src/ops/array_ops.ts +++ b/tfjs-core/src/ops/array_ops.ts @@ -237,7 +237,7 @@ function spaceToBatchND_( blockShape.toString()}`); const grad = (dy: T) => { - return {$x: () => dy.batchToSpaceND(blockShape, paddings)}; + return {$x: () => dy.batchToSpaceND(blockShape, paddings) as T}; }; return ENGINE.runKernelFunc( diff --git a/tfjs-core/src/tests.ts b/tfjs-core/src/tests.ts index 5cc01c14ee7..2b8348169cc 100644 --- a/tfjs-core/src/tests.ts +++ b/tfjs-core/src/tests.ts @@ -43,6 +43,7 @@ import './ops/add_test'; import './ops/arithmetic_test'; import './ops/array_ops_test'; import './ops/axis_util_test'; +import './ops/batch_to_space_nd_test'; import './ops/batchnorm_test'; import './ops/binary_ops_test'; import './ops/boolean_mask_test'; From facfcd83d4f4915a81bbc3081e77d3dc5b09e3a8 Mon Sep 17 00:00:00 2001 From: Ann Yuan Date: Thu, 14 May 2020 16:08:31 -0400 Subject: [PATCH 4/4] clean --- tfjs-core/src/gradients/BatchToSpaceND_grad.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tfjs-core/src/gradients/BatchToSpaceND_grad.ts b/tfjs-core/src/gradients/BatchToSpaceND_grad.ts index ccf57c96849..68e5bde2557 100644 --- a/tfjs-core/src/gradients/BatchToSpaceND_grad.ts +++ b/tfjs-core/src/gradients/BatchToSpaceND_grad.ts @@ -23,9 +23,7 @@ import {Tensor} from '../tensor'; export const batchToSpaceNDGradConfig: GradConfig = { kernelName: BatchToSpaceND, gradFunc: (dy: Tensor, saved: Tensor[], attrs: NamedAttrMap) => { - const batchToSpaceNDAttrs: BatchToSpaceNDAttrs = - attrs as {} as BatchToSpaceNDAttrs; - const {blockShape, crops} = batchToSpaceNDAttrs; + const {blockShape, crops} = attrs as {} as BatchToSpaceNDAttrs; return {x: () => spaceToBatchND(dy, blockShape, crops)}; } };