From a74bfa9e7c7673477d9166ba5b0eb271562d1b05 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 30 Mar 2020 12:29:44 -0700 Subject: [PATCH 1/8] Modularize batchnorm. --- tfjs-core/src/ops/batchnorm.ts | 299 +----- tfjs-core/src/ops/batchnorm2d.ts | 95 ++ tfjs-core/src/ops/batchnorm2d_test.ts | 458 +++++++++ tfjs-core/src/ops/batchnorm3d.ts | 94 ++ tfjs-core/src/ops/batchnorm3d_test.ts | 296 ++++++ tfjs-core/src/ops/batchnorm4d.ts | 93 ++ tfjs-core/src/ops/batchnorm4d_test.ts | 233 +++++ tfjs-core/src/ops/batchnorm_test.ts | 950 ------------------ tfjs-core/src/ops/batchnorm_util.ts | 24 + tfjs-core/src/ops/ops.ts | 5 +- tfjs-core/src/public/chained_ops/batchnorm.ts | 40 + .../chained_ops/register_all_chained_ops.ts | 1 + .../register_all_chained_ops_test.ts | 4 +- tfjs-core/src/tensor.ts | 17 - tfjs-core/src/tests.ts | 4 +- 15 files changed, 1384 insertions(+), 1229 deletions(-) create mode 100644 tfjs-core/src/ops/batchnorm2d.ts create mode 100644 tfjs-core/src/ops/batchnorm2d_test.ts create mode 100644 tfjs-core/src/ops/batchnorm3d.ts create mode 100644 tfjs-core/src/ops/batchnorm3d_test.ts create mode 100644 tfjs-core/src/ops/batchnorm4d.ts create mode 100644 tfjs-core/src/ops/batchnorm4d_test.ts delete mode 100644 tfjs-core/src/ops/batchnorm_test.ts create mode 100644 tfjs-core/src/ops/batchnorm_util.ts create mode 100644 tfjs-core/src/public/chained_ops/batchnorm.ts diff --git a/tfjs-core/src/ops/batchnorm.ts b/tfjs-core/src/ops/batchnorm.ts index 94d58fb79a4..317da3a579c 100644 --- a/tfjs-core/src/ops/batchnorm.ts +++ b/tfjs-core/src/ops/batchnorm.ts @@ -1,6 +1,6 @@ /** * @license - * Copyright 2018 Google Inc. All Rights Reserved. + * 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 @@ -16,8 +16,7 @@ */ import {ENGINE} from '../engine'; -import {deprecationWarn} from '../globals'; -import {Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D} from '../tensor'; +import {Tensor, Tensor1D, Tensor4D} from '../tensor'; import {convertToTensor} from '../tensor_util_env'; import {Rank, ShapeMap, TensorLike} from '../types'; import * as util from '../util'; @@ -27,180 +26,17 @@ import {op} from './operation'; import {scalar} from './tensor_ops'; import {tile} from './tile'; import {rsqrt} from './unary_ops'; - -/** - * Batch normalization, strictly for 2D. For the more relaxed version, see - * `tf.batchNorm`. - * - * @param x The input Tensor. - * @param mean A mean Tensor. - * @param variance A variance Tensor. - * @param offset An offset Tensor. - * @param scale A scale Tensor. - * @param varianceEpsilon A small float number to avoid dividing by 0. - */ -function batchNorm2d_( - x: Tensor2D|TensorLike, mean: Tensor2D|Tensor1D|TensorLike, - variance: Tensor2D|Tensor1D|TensorLike, - offset?: Tensor2D|Tensor1D|TensorLike, scale?: Tensor2D|Tensor1D|TensorLike, - varianceEpsilon?: number): Tensor2D { - const $x = convertToTensor(x, 'x', 'batchNorm'); - const $mean = convertToTensor(mean, 'mean', 'batchNorm'); - const $variance = convertToTensor(variance, 'variance', 'batchNorm'); - let $scale: Tensor2D|Tensor1D; - if (scale != null) { - $scale = convertToTensor(scale, 'scale', 'batchNorm'); - } - let $offset: Tensor2D|Tensor1D; - if (offset != null) { - $offset = convertToTensor(offset, 'offset', 'batchNorm'); - } - util.assert( - $x.rank === 2, - () => `Error in batchNorm3D: x must be rank 3 but got rank ` + - `${$x.rank}.`); - util.assert( - $mean.rank === 2 || $mean.rank === 1, - () => `Error in batchNorm2D: mean must be rank 2 or rank 1 but ` + - `got rank ${$mean.rank}.`); - util.assert( - $variance.rank === 2 || $variance.rank === 1, - () => `Error in batchNorm2D: variance must be rank 2 or rank 1 ` + - `but got rank ${$variance.rank}.`); - if ($scale != null) { - util.assert( - $scale.rank === 2 || $scale.rank === 1, - () => `Error in batchNorm2D: scale must be rank 2 or rank 1 ` + - `but got rank ${$scale.rank}.`); - } - if ($offset != null) { - util.assert( - $offset.rank === 2 || $offset.rank === 1, - () => `Error in batchNorm2D: offset must be rank 2 or rank 1 ` + - `but got rank ${$offset.rank}.`); - } - - return batchNorm_($x, $mean, $variance, $offset, $scale, varianceEpsilon); -} - -/** - * Batch normalization, strictly for 3D. For the more relaxed version, see - * `tf.batchNorm`. - * - * @param x The input Tensor. - * @param mean A mean Tensor. - * @param variance A variance Tensor. - * @param offset An offset Tensor. - * @param scale A scale Tensor. - * @param varianceEpsilon A small float number to avoid dividing by 0. - */ -function batchNorm3d_( - x: Tensor3D|TensorLike, mean: Tensor3D|Tensor1D|TensorLike, - variance: Tensor3D|Tensor1D|TensorLike, - offset?: Tensor3D|Tensor1D|TensorLike, scale?: Tensor3D|Tensor1D|TensorLike, - varianceEpsilon?: number): Tensor3D { - const $x = convertToTensor(x, 'x', 'batchNorm'); - const $mean = convertToTensor(mean, 'mean', 'batchNorm'); - const $variance = convertToTensor(variance, 'variance', 'batchNorm'); - let $scale: Tensor3D|Tensor1D; - if (scale != null) { - $scale = convertToTensor(scale, 'scale', 'batchNorm'); - } - let $offset: Tensor3D|Tensor1D; - if (offset != null) { - $offset = convertToTensor(offset, 'offset', 'batchNorm'); - } - util.assert( - $x.rank === 3, - () => `Error in batchNorm3D: x must be rank 3 but got rank ` + - `${$x.rank}.`); - util.assert( - $mean.rank === 3 || $mean.rank === 1, - () => `Error in batchNorm3D: mean must be rank 3 or rank 1 but ` + - `got rank ${$mean.rank}.`); - util.assert( - $variance.rank === 3 || $variance.rank === 1, - () => `Error in batchNorm3D: variance must be rank 3 or rank 1 ` + - `but got rank ${$variance.rank}.`); - if ($scale != null) { - util.assert( - $scale.rank === 3 || $scale.rank === 1, - () => `Error in batchNorm3D: scale must be rank 3 or rank 1 ` + - `but got rank ${$scale.rank}.`); - } - if ($offset != null) { - util.assert( - $offset.rank === 3 || $offset.rank === 1, - () => `Error in batchNorm3D: offset must be rank 3 or rank 1 ` + - `but got rank ${$offset.rank}.`); - } - - return batchNorm_($x, $mean, $variance, $offset, $scale, varianceEpsilon); -} - -/** - * Batch normalization, strictly for 4D. For the more relaxed version, see - * `tf.batchNorm`. - * - * @param x The input Tensor. - * @param mean A mean Tensor. - * @param variance A variance Tensor. - * @param offset An offset Tensor. - * @param scale A scale Tensor. - * @param varianceEpsilon A small float number to avoid dividing by 0. - */ -function batchNorm4d_( - x: Tensor4D|TensorLike, mean: Tensor4D|Tensor1D|TensorLike, - variance: Tensor4D|Tensor1D|TensorLike, - offset?: Tensor4D|Tensor1D|TensorLike, scale?: Tensor4D|Tensor1D|TensorLike, - varianceEpsilon?: number): Tensor4D { - const $x = convertToTensor(x, 'x', 'batchNorm'); - const $mean = convertToTensor(mean, 'mean', 'batchNorm'); - const $variance = convertToTensor(variance, 'variance', 'batchNorm'); - let $scale: Tensor4D|Tensor1D; - if (scale != null) { - $scale = convertToTensor(scale, 'scale', 'batchNorm'); - } - let $offset: Tensor4D|Tensor1D; - if (offset != null) { - $offset = convertToTensor(offset, 'offset', 'batchNorm'); - } - util.assert( - $x.rank === 4, - () => `Error in batchNorm4D: x must be rank 4 but got rank ` + - `${$x.rank}.`); - util.assert( - $mean.rank === 4 || $mean.rank === 1, - () => `Error in batchNorm4D: mean must be rank 4 or rank 1 but ` + - `got rank ${$mean.rank}.`); - util.assert( - $variance.rank === 4 || $variance.rank === 1, - () => `Error in batchNorm4D: variance must be rank 4 or rank 1 ` + - `but got rank ${$variance.rank}.`); - if ($scale != null) { - util.assert( - $scale.rank === 4 || $scale.rank === 1, - () => `Error in batchNorm4D: scale must be rank 4 or rank 1 ` + - `but got rank ${$scale.rank}.`); - } - if ($offset != null) { - util.assert( - $offset.rank === 4 || $offset.rank === 1, - () => `Error in batchNorm4D: offset must be rank 4 or rank 1 ` + - `but got rank ${$offset.rank}.`); - } - return batchNorm_($x, $mean, $variance, $offset, $scale, varianceEpsilon); -} +import {warnDeprecation} from './batchnorm_util'; /** * @deprecated Please use `tf.batchNorm` instead and note the positional * argument change of scale, offset, and varianceEpsilon. */ function batchNormalization_( - x: Tensor|TensorLike, mean: Tensor|Tensor1D|TensorLike, - variance: Tensor|Tensor1D|TensorLike, varianceEpsilon = .001, - scale?: Tensor|Tensor1D|TensorLike, - offset?: Tensor|Tensor1D|TensorLike): Tensor { + x: Tensor | TensorLike, mean: Tensor | Tensor1D | TensorLike, + variance: Tensor | Tensor1D | TensorLike, varianceEpsilon = .001, + scale?: Tensor | Tensor1D | TensorLike, + offset?: Tensor | Tensor1D | TensorLike): Tensor { warnDeprecation(); return batchNorm_(x, mean, variance, offset, scale, varianceEpsilon); } @@ -231,38 +67,38 @@ function batchNormalization_( */ /** @doc {heading: 'Operations', subheading: 'Normalization'} */ function batchNorm_( - x: Tensor|TensorLike, mean: Tensor|Tensor1D|TensorLike, - variance: Tensor|Tensor1D|TensorLike, - offset?: Tensor|Tensor1D|TensorLike, - scale?: Tensor|Tensor1D|TensorLike, - varianceEpsilon?: number): Tensor { + x: Tensor | TensorLike, mean: Tensor | Tensor1D | TensorLike, + variance: Tensor | Tensor1D | TensorLike, + offset?: Tensor | Tensor1D | TensorLike, + scale?: Tensor | Tensor1D | TensorLike, + varianceEpsilon?: number): Tensor { if (varianceEpsilon == null) { varianceEpsilon = 0.001; } const $x = convertToTensor(x, 'x', 'batchNorm'); const $mean = convertToTensor(mean, 'mean', 'batchNorm'); const $variance = convertToTensor(variance, 'variance', 'batchNorm'); - let $scale: Tensor|Tensor1D; + let $scale: Tensor | Tensor1D; if (scale != null) { $scale = convertToTensor(scale, 'scale', 'batchNorm'); } - let $offset: Tensor|Tensor1D; + let $offset: Tensor | Tensor1D; if (offset != null) { $offset = convertToTensor(offset, 'offset', 'batchNorm'); } util.assert( - $mean.rank === $variance.rank, - () => 'Batch normalization gradient requires mean and variance to have ' + - 'equal ranks.'); + $mean.rank === $variance.rank, + () => 'Batch normalization gradient requires mean and variance to have ' + + 'equal ranks.'); util.assert( - $offset == null || $mean.rank === $offset.rank, - () => 'Batch normalization gradient requires mean and offset to have ' + - 'equal ranks.'); + $offset == null || $mean.rank === $offset.rank, + () => 'Batch normalization gradient requires mean and offset to have ' + + 'equal ranks.'); util.assert( - $scale == null || $mean.rank === $scale.rank, - () => 'Batch normalization gradient requires mean and scale to have ' + - 'equal ranks.'); + $scale == null || $mean.rank === $scale.rank, + () => 'Batch normalization gradient requires mean and scale to have ' + + 'equal ranks.'); let x4D: Tensor4D; if ($x.rank === 0 || $x.rank === 1) { @@ -277,7 +113,7 @@ function batchNorm_( const der = (dy: Tensor, saved: Tensor[]) => { type Saved = [ - Tensor, Tensor| Tensor1D, Tensor| Tensor1D, Tensor| Tensor1D + Tensor, Tensor | Tensor1D, Tensor | Tensor1D, Tensor | Tensor1D ]; const [$x, $mean, $variance, $scale] = saved as Saved; const scaleValue = $scale == null ? scalar(1) : $scale; @@ -294,16 +130,16 @@ function batchNorm_( const dyTimesScaleValue = dy.mul(scaleValue); const oneOverSqrtVariance = rsqrt($variance.add(scalar(varianceEpsilon))); const minusHalfRCube = oneOverSqrtVariance.mul(oneOverSqrtVariance) - .mul(oneOverSqrtVariance) - .mul(scalar(-0.5)); + .mul(oneOverSqrtVariance) + .mul(scalar(-0.5)); const derX = () => { if ($mean.rank === 1) { return dy - .mul(tile( - oneOverSqrtVariance.as4D(1, 1, 1, $mean.shape[0]), tileShape)) - .mul(scaleValue) - .reshape($x.shape); + .mul(tile( + oneOverSqrtVariance.as4D(1, 1, 1, $mean.shape[0]), tileShape)) + .mul(scaleValue) + .reshape($x.shape); } else { return dy.mul(oneOverSqrtVariance).mul(scaleValue).reshape($x.shape); } @@ -349,20 +185,20 @@ function batchNorm_( const inputsToSave = [$x, $mean, $variance, $scale]; const res = ENGINE.runKernelFunc( - (backend, save) => { - const res = backend.batchNormalization( - x4D, batchnormReshape4D($mean), batchnormReshape4D($variance), - varianceEpsilon, batchnormReshape4D($scale), - batchnormReshape4D($offset)); - save([$x, $mean, $variance, $scale]); - return res; - }, - {x: $x, mean: $mean, variance: $variance, scale: $scale, offset: $offset}, - der, 'BatchNormalization', {varianceEpsilon}, inputsToSave); + (backend, save) => { + const res = backend.batchNormalization( + x4D, batchnormReshape4D($mean), batchnormReshape4D($variance), + varianceEpsilon, batchnormReshape4D($scale), + batchnormReshape4D($offset)); + save([$x, $mean, $variance, $scale]); + return res; + }, + {x: $x, mean: $mean, variance: $variance, scale: $scale, offset: $offset}, + der, 'BatchNormalization', {varianceEpsilon}, inputsToSave); return res.reshape($x.shape); } -function batchnormReshape4D(x: Tensor): Tensor4D|Tensor1D { +function batchnormReshape4D(x: Tensor): Tensor4D | Tensor1D { if (x == null) { return null; } @@ -378,58 +214,5 @@ function batchnormReshape4D(x: Tensor): Tensor4D|Tensor1D { return x as Tensor4D; } -/** - * @deprecated Please use `tf.batchNorm2d` instead and note the positional - * argument change of scale, offset, and varianceEpsilon. - */ -function batchNormalization2d_( - x: Tensor2D|TensorLike, mean: Tensor2D|Tensor1D|TensorLike, - variance: Tensor2D|Tensor1D|TensorLike, varianceEpsilon = .001, - scale?: Tensor2D|Tensor1D|TensorLike, - offset?: Tensor2D|Tensor1D|TensorLike): Tensor2D { - warnDeprecation(); - return batchNorm2d_(x, mean, variance, offset, scale, varianceEpsilon); -} - -/** - * @deprecated Please use `tf.batchNorm3d` instead and note the positional - * argument change of scale, offset, and varianceEpsilon. - */ -function batchNormalization3d_( - x: Tensor3D|TensorLike, mean: Tensor3D|Tensor1D|TensorLike, - variance: Tensor3D|Tensor1D|TensorLike, varianceEpsilon = .001, - scale?: Tensor3D|Tensor1D|TensorLike, - offset?: Tensor3D|Tensor1D|TensorLike): Tensor3D { - warnDeprecation(); - return batchNorm3d_(x, mean, variance, offset, scale, varianceEpsilon); -} - -/** - * @deprecated Please use `tf.batchNorm4d` instead and note the positional - * argument change of scale, offset, and varianceEpsilon. - */ -function batchNormalization4d_( - x: Tensor4D|TensorLike, mean: Tensor4D|Tensor1D|TensorLike, - variance: Tensor4D|Tensor1D|TensorLike, varianceEpsilon = .001, - scale?: Tensor4D|Tensor1D|TensorLike, - offset?: Tensor4D|Tensor1D|TensorLike): Tensor4D { - warnDeprecation(); - return batchNorm4d_(x, mean, variance, offset, scale, varianceEpsilon); -} - -function warnDeprecation() { - deprecationWarn( - 'tf.batchNormalization() is going away. ' + - 'Use tf.batchNorm() instead, and note the positional argument change ' + - 'of scale, offset, and varianceEpsilon'); -} - -export const batchNormalization2d = op({batchNormalization2d_}); -export const batchNormalization3d = op({batchNormalization3d_}); -export const batchNormalization4d = op({batchNormalization4d_}); export const batchNormalization = op({batchNormalization_}); - export const batchNorm = op({batchNorm_}); -export const batchNorm2d = op({batchNorm2d_}); -export const batchNorm3d = op({batchNorm3d_}); -export const batchNorm4d = op({batchNorm4d_}); diff --git a/tfjs-core/src/ops/batchnorm2d.ts b/tfjs-core/src/ops/batchnorm2d.ts new file mode 100644 index 00000000000..7f8fe25ae84 --- /dev/null +++ b/tfjs-core/src/ops/batchnorm2d.ts @@ -0,0 +1,95 @@ +/** + * @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 {Tensor1D, Tensor2D} from '../tensor'; +import {convertToTensor} from '../tensor_util_env'; +import {TensorLike} from '../types'; +import * as util from '../util'; + +import {op} from './operation'; +import {batchNorm} from './batchnorm'; +import {warnDeprecation} from './batchnorm_util'; + +/** + * Batch normalization, strictly for 2D. For the more relaxed version, see + * `tf.batchNorm`. + * + * @param x The input Tensor. + * @param mean A mean Tensor. + * @param variance A variance Tensor. + * @param offset An offset Tensor. + * @param scale A scale Tensor. + * @param varianceEpsilon A small float number to avoid dividing by 0. + */ +function batchNorm2d_( + x: Tensor2D | TensorLike, mean: Tensor2D | Tensor1D | TensorLike, + variance: Tensor2D | Tensor1D | TensorLike, + offset?: Tensor2D | Tensor1D | TensorLike, scale?: Tensor2D | Tensor1D | TensorLike, + varianceEpsilon?: number): Tensor2D { + const $x = convertToTensor(x, 'x', 'batchNorm'); + const $mean = convertToTensor(mean, 'mean', 'batchNorm'); + const $variance = convertToTensor(variance, 'variance', 'batchNorm'); + let $scale: Tensor2D | Tensor1D; + if (scale != null) { + $scale = convertToTensor(scale, 'scale', 'batchNorm'); + } + let $offset: Tensor2D | Tensor1D; + if (offset != null) { + $offset = convertToTensor(offset, 'offset', 'batchNorm'); + } + util.assert( + $x.rank === 2, + () => `Error in batchNorm3D: x must be rank 3 but got rank ` + + `${$x.rank}.`); + util.assert( + $mean.rank === 2 || $mean.rank === 1, + () => `Error in batchNorm2D: mean must be rank 2 or rank 1 but ` + + `got rank ${$mean.rank}.`); + util.assert( + $variance.rank === 2 || $variance.rank === 1, + () => `Error in batchNorm2D: variance must be rank 2 or rank 1 ` + + `but got rank ${$variance.rank}.`); + if ($scale != null) { + util.assert( + $scale.rank === 2 || $scale.rank === 1, + () => `Error in batchNorm2D: scale must be rank 2 or rank 1 ` + + `but got rank ${$scale.rank}.`); + } + if ($offset != null) { + util.assert( + $offset.rank === 2 || $offset.rank === 1, + () => `Error in batchNorm2D: offset must be rank 2 or rank 1 ` + + `but got rank ${$offset.rank}.`); + } + + return batchNorm($x, $mean, $variance, $offset, $scale, varianceEpsilon); +} + +/** + * @deprecated Please use `tf.batchNorm2d` instead and note the positional + * argument change of scale, offset, and varianceEpsilon. + */ +function batchNormalization2d_( + x: Tensor2D | TensorLike, mean: Tensor2D | Tensor1D | TensorLike, + variance: Tensor2D | Tensor1D | TensorLike, varianceEpsilon = .001, + scale?: Tensor2D | Tensor1D | TensorLike, + offset?: Tensor2D | Tensor1D | TensorLike): Tensor2D { + warnDeprecation(); + return batchNorm2d_(x, mean, variance, offset, scale, varianceEpsilon); +} + +export const batchNormalization2d = op({batchNormalization2d_}); +export const batchNorm2d = op({batchNorm2d_}); diff --git a/tfjs-core/src/ops/batchnorm2d_test.ts b/tfjs-core/src/ops/batchnorm2d_test.ts new file mode 100644 index 00000000000..2eb33fe0d64 --- /dev/null +++ b/tfjs-core/src/ops/batchnorm2d_test.ts @@ -0,0 +1,458 @@ +/** + * @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 * as tf from '../index'; +import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; +import {expectArraysClose} from '../test_util'; + +describeWithFlags('batchNorm2D', ALL_ENVS, () => { + it('simple batchnorm2D, no offset or scale, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const varianceEpsilon = .001; + + const result = tf.batchNorm2d( + xT, meanT, varianceT, undefined, undefined, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + expectArraysClose(await result.data(), [ + (x.get(0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + it('simple batchnorm2D, no offset, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const scaleT = tf.tensor1d([4, 5]); + const varianceEpsilon = .001; + + const result = tf.batchNorm2d( + xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const scale = await scaleT.buffer(); + expectArraysClose(await result.data(), [ + (x.get(0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm2D, no scale, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = tf.batchNorm2d( + xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); + + const offset = await offsetT.array(); + const mean = await meanT.array(); + const variance = await varianceT.array(); + const x = await xT.array(); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('simple batchnorm2D, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([3, 4]); + const scaleT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = + tf.batchNorm2d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); + + const offset = await offsetT.array(); + const mean = await meanT.array(); + const variance = await varianceT.array(); + const scale = await scaleT.array(); + const x = await xT.array(); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('simple batchnorm2D gradients, 2x2', async () => { + const x = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const offset = tf.tensor1d([3, 4]); + const scale = tf.tensor1d([2, 5]); + + const varianceEpsilon = .001; + + const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); + const [gradX, gradMean, gradVariance, gradOffset, gradScale] = tf.grads( + (x: tf.Tensor2D, mean: tf.Tensor1D, variance: tf.Tensor1D, + offset: tf.Tensor1D, scale: tf.Tensor1D) => + tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon))( + [x, mean, variance, offset, scale], dy); + expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); + expect(gradX.shape).toEqual([2, 2]); + expectArraysClose(await gradMean.data(), [-2.828, -5.773]); + expect(gradMean.shape).toEqual([2]); + expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); + expect(gradVariance.shape).toEqual([2]); + expectArraysClose(await gradOffset.data(), [2, 2]); + expect(gradOffset.shape).toEqual([2]); + expectArraysClose(await gradScale.data(), [6.362, 13.277]); + expect(gradScale.shape).toEqual([2]); + }); + + it('gradient with clones batchnorm2D', async () => { + const x = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const offset = tf.tensor1d([3, 4]); + const scale = tf.tensor1d([2, 5]); + + const varianceEpsilon = .001; + + const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); + const [gradX, gradMean, gradVariance, gradOffset, gradScale] = tf.grads( + (x: tf.Tensor2D, mean: tf.Tensor1D, variance: tf.Tensor1D, + offset: tf.Tensor1D, scale: tf.Tensor1D) => + tf.batchNorm2d( + x.clone(), mean.clone(), variance.clone(), offset.clone(), + scale.clone(), varianceEpsilon) + .clone())([x, mean, variance, offset, scale], dy); + expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); + expect(gradX.shape).toEqual([2, 2]); + expectArraysClose(await gradMean.data(), [-2.828, -5.773]); + expect(gradMean.shape).toEqual([2]); + expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); + expect(gradVariance.shape).toEqual([2]); + expectArraysClose(await gradOffset.data(), [2, 2]); + expect(gradOffset.shape).toEqual([2]); + expectArraysClose(await gradScale.data(), [6.362, 13.277]); + expect(gradScale.shape).toEqual([2]); + }); + + it('batchnorm2D gradients, same shapes in x, mean and variance', async () => { + const x = tf.tensor2d([10, 20, 30, 40], [2, 2]); + const mean = tf.tensor2d([0, 5, 10, 15], [2, 2]); + const variance = tf.tensor2d([2, 4, 6, 8], [2, 2]); + const scale = tf.tensor2d([2, 5, 2, 5], [2, 2]); + const offset = tf.tensor2d([0, 0, 0, 0], [2, 2]); + + const varianceEpsilon = .001; + + const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); + const gradX = tf.grad( + (x: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [1.414, 2.500, 0.816, 1.768]); + expect(gradX.shape).toEqual([2, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [-1.414, -2.500, -0.816, -1.768]); + expect(gradMean.shape).toEqual([2, 2]); + const gradVariance = tf.grad( + (variance: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose( + await gradVariance.data(), [-3.533, -4.686, -1.360, -2.762]); + expect(gradVariance.shape).toEqual([2, 2]); + const gradOffset = tf.grad( + (offset: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), [1, 1, 1, 1]); + expect(gradOffset.shape).toEqual([2, 2]); + const gradScale = tf.grad( + (scale: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [7.069, 7.499, 8.164, 8.838]); + expect(gradScale.shape).toEqual([2, 2]); + }); + + it('gradient with clones', () => { + const x = tf.zeros([2, 2]); + const mean = tf.zeros([2, 2]); + const variance = tf.zeros([2, 2]); + const scale = tf.zeros([2, 2]); + const offset = tf.zeros([2, 2]); + + const varianceEpsilon = .001; + + const gradF = tf.grads( + (x: tf.Tensor2D, mean: tf.Tensor2D, variance: tf.Tensor2D, + offset: tf.Tensor2D, scale: tf.Tensor2D) => + tf.batchNorm2d( + x.clone(), mean.clone(), variance.clone(), offset.clone(), + scale.clone(), varianceEpsilon) + .clone()); + const [gradX, gradMean, gradVariance, gradOffset, gradScale] = + gradF([x, mean, variance, offset, scale]); + expect(gradX.shape).toEqual(x.shape); + expect(gradMean.shape).toEqual(mean.shape); + expect(gradVariance.shape).toEqual(variance.shape); + expect(gradOffset.shape).toEqual(offset.shape); + expect(gradScale.shape).toEqual(scale.shape); + }); + + it('batchnorm2D matches tensorflow, 3x3', async () => { + const x = tf.tensor2d( + [ + 0.3136892, 0.92389025, 0.594782, 0.05021042, 0.67545404, 0.93910035, + 0.13277993, 0.96474269, 0.88608916 + ], + [3, 3]); + const mean = tf.tensor1d([0.19526312, 0.74857256, 0.45166398]); + const variance = tf.tensor1d([0.22963001, 0.61521992, 0.46623685]); + const offset = tf.tensor1d([0.43098484, 0.77712237, 0.47916298]); + const scale = tf.tensor1d([0.62186907, 0.85673736, 0.19201061]); + const varianceEpsilon = .001; + + const result = + tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + 0.58433646, 0.96846228, 0.51936529, 0.24315402, 0.69732157, 0.61608542, + 0.35007446, 1.01304821, 0.60119441 + ]); + }); + + it('throws when passed x as a non-tensor', () => { + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + + expect(() => tf.batchNorm({} as tf.Tensor, mean, variance)) + .toThrowError(/Argument 'x' passed to 'batchNorm' must be a Tensor/); + }); + it('throws when passed mean as a non-tensor', () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const variance = tf.tensor1d([2, 3]); + + expect(() => tf.batchNorm(x, {} as tf.Tensor, variance)) + .toThrowError(/Argument 'mean' passed to 'batchNorm' must be a Tensor/); + }); + it('throws when passed variance as a non-tensor', () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const mean = tf.tensor1d([1, 2]); + + const e = /Argument 'variance' passed to 'batchNorm' must be a Tensor/; + expect(() => tf.batchNorm(x, mean, {} as tf.Tensor)).toThrowError(e); + }); + it('throws when passed scale as a non-tensor', () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const epsilon = .001; + + expect(() => tf.batchNorm(x, mean, variance, epsilon, {} as tf.Tensor)) + .toThrowError( + /Argument 'scale' passed to 'batchNorm' must be a Tensor/); + }); + it('throws when passed offset as a non-tensor', () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const epsilon = .001; + const scale = tf.tensor1d([0.62186907, 0.85673736, 0.19201061]); + + const e = /Argument 'offset' passed to 'batchNorm' must be a Tensor/; + expect( + () => tf.batchNorm(x, mean, variance, {} as tf.Tensor, scale, epsilon)) + .toThrowError(e); + }); + + it('accepts a tensor-like object', async () => { + const x = [[2, 4], [9, 23]]; + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const result = + tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('throws error when x is a string tensor', () => { + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const f = () => tf.batchNorm2d( + [['a', 'b'], ['c', 'd']], mean, variance, offset, scale, + varianceEpsilon); + expect(f).toThrowError( + /Argument 'x' passed to 'batchNorm' must be numeric/); + }); + + it('throws error when mean is a string tensor', () => { + const x = [[2, 4], [9, 23]]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const f = () => + tf.batchNorm2d(x, ['a', 'b'], variance, offset, scale, varianceEpsilon); + expect(f).toThrowError( + /Argument 'mean' passed to 'batchNorm' must be numeric/); + }); + + it('throws error when variance is a string tensor', () => { + const x = [[2, 4], [9, 23]]; + const mean = [1, 2]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const f = () => + tf.batchNorm2d(x, mean, ['a', 'b'], offset, scale, varianceEpsilon); + expect(f).toThrowError(/'variance' passed to 'batchNorm' must be numeric/); + }); + + it('throws error when scale is a string tensor', () => { + const x = [[2, 4], [9, 23]]; + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + + const varianceEpsilon = .001; + + const f = () => + tf.batchNorm2d(x, mean, variance, offset, ['a', 'b'], varianceEpsilon); + expect(f).toThrowError(/'scale' passed to 'batchNorm' must be numeric/); + }); + + it('throws error when offset is a string tensor', () => { + const x = [[2, 4], [9, 23]]; + const mean = [1, 2]; + const variance = [2, 3]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const f = () => + tf.batchNorm2d(x, mean, variance, ['a', 'b'], scale, varianceEpsilon); + expect(f).toThrowError(/'offset' passed to 'batchNorm' must be numeric/); + }); +}); + +describeWithFlags('deprecated batchNormalization', ALL_ENVS, () => { + beforeAll(() => { + // Silence deprecation warnings. + spyOn(console, 'warn'); + }); + + it('simple batchnorm2D, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([3, 4]); + const scaleT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = tf.batchNormalization( + xT, meanT, varianceT, varianceEpsilon, scaleT, offsetT); + + const offset = await offsetT.array(); + const mean = await meanT.array(); + const variance = await varianceT.array(); + const scale = await scaleT.array(); + const x = await xT.array(); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + + const result2 = tf.batchNormalization2d( + xT, meanT, varianceT, varianceEpsilon, scaleT, offsetT); + expectArraysClose(await result.data(), await result2.data()); + }); +}); diff --git a/tfjs-core/src/ops/batchnorm3d.ts b/tfjs-core/src/ops/batchnorm3d.ts new file mode 100644 index 00000000000..64638a4311e --- /dev/null +++ b/tfjs-core/src/ops/batchnorm3d.ts @@ -0,0 +1,94 @@ +/** + * @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 {Tensor1D, Tensor3D} from '../tensor'; +import {convertToTensor} from '../tensor_util_env'; +import {TensorLike} from '../types'; +import * as util from '../util'; + +import {op} from './operation'; +import {batchNorm} from './batchnorm'; +import {warnDeprecation} from './batchnorm_util'; +/** + * Batch normalization, strictly for 3D. For the more relaxed version, see + * `tf.batchNorm`. + * + * @param x The input Tensor. + * @param mean A mean Tensor. + * @param variance A variance Tensor. + * @param offset An offset Tensor. + * @param scale A scale Tensor. + * @param varianceEpsilon A small float number to avoid dividing by 0. + */ +function batchNorm3d_( + x: Tensor3D | TensorLike, mean: Tensor3D | Tensor1D | TensorLike, + variance: Tensor3D | Tensor1D | TensorLike, + offset?: Tensor3D | Tensor1D | TensorLike, scale?: Tensor3D | Tensor1D | TensorLike, + varianceEpsilon?: number): Tensor3D { + const $x = convertToTensor(x, 'x', 'batchNorm'); + const $mean = convertToTensor(mean, 'mean', 'batchNorm'); + const $variance = convertToTensor(variance, 'variance', 'batchNorm'); + let $scale: Tensor3D | Tensor1D; + if (scale != null) { + $scale = convertToTensor(scale, 'scale', 'batchNorm'); + } + let $offset: Tensor3D | Tensor1D; + if (offset != null) { + $offset = convertToTensor(offset, 'offset', 'batchNorm'); + } + util.assert( + $x.rank === 3, + () => `Error in batchNorm3D: x must be rank 3 but got rank ` + + `${$x.rank}.`); + util.assert( + $mean.rank === 3 || $mean.rank === 1, + () => `Error in batchNorm3D: mean must be rank 3 or rank 1 but ` + + `got rank ${$mean.rank}.`); + util.assert( + $variance.rank === 3 || $variance.rank === 1, + () => `Error in batchNorm3D: variance must be rank 3 or rank 1 ` + + `but got rank ${$variance.rank}.`); + if ($scale != null) { + util.assert( + $scale.rank === 3 || $scale.rank === 1, + () => `Error in batchNorm3D: scale must be rank 3 or rank 1 ` + + `but got rank ${$scale.rank}.`); + } + if ($offset != null) { + util.assert( + $offset.rank === 3 || $offset.rank === 1, + () => `Error in batchNorm3D: offset must be rank 3 or rank 1 ` + + `but got rank ${$offset.rank}.`); + } + + return batchNorm($x, $mean, $variance, $offset, $scale, varianceEpsilon); +} + +/** + * @deprecated Please use `tf.batchNorm3d` instead and note the positional + * argument change of scale, offset, and varianceEpsilon. + */ +function batchNormalization3d_( + x: Tensor3D | TensorLike, mean: Tensor3D | Tensor1D | TensorLike, + variance: Tensor3D | Tensor1D | TensorLike, varianceEpsilon = .001, + scale?: Tensor3D | Tensor1D | TensorLike, + offset?: Tensor3D | Tensor1D | TensorLike): Tensor3D { + warnDeprecation(); + return batchNorm3d_(x, mean, variance, offset, scale, varianceEpsilon); +} + +export const batchNormalization3d = op({batchNormalization3d_}); +export const batchNorm3d = op({batchNorm3d_}); diff --git a/tfjs-core/src/ops/batchnorm3d_test.ts b/tfjs-core/src/ops/batchnorm3d_test.ts new file mode 100644 index 00000000000..02473575604 --- /dev/null +++ b/tfjs-core/src/ops/batchnorm3d_test.ts @@ -0,0 +1,296 @@ +/** + * @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 * as tf from '../index'; +import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; +import {expectArraysClose} from '../test_util'; + +describeWithFlags('batchNorm3D', ALL_ENVS, () => { + it('simple batchnorm3D, no offset or scale, 2x1x2', async () => { + const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const varianceEpsilon = .001; + + const result = tf.batchNorm3d( + xT, meanT, varianceT, undefined, undefined, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + expectArraysClose(await result.data(), [ + (x.get(0, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm3D, no offset, 2x1x2', async () => { + const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const scaleT = tf.tensor1d([4, 5]); + const varianceEpsilon = .001; + + const result = tf.batchNorm3d( + xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const scale = await scaleT.buffer(); + expectArraysClose(await result.data(), [ + (x.get(0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm3D, no scale, 2x1x2', async () => { + const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = tf.batchNorm3d( + xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const offset = await offsetT.buffer(); + expectArraysClose(await result.data(), [ + offset.get(0) + + (x.get(0, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(0, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon), + offset.get(0) + + (x.get(1, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(1, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm3D, 2x1x2', async () => { + const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([3, 4]); + const scaleT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = + tf.batchNorm3d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const offset = await offsetT.buffer(); + const scale = await scaleT.buffer(); + + expectArraysClose(await result.data(), [ + offset.get(0) + + (x.get(0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + offset.get(0) + + (x.get(1, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(1, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('accepts a tensor-like object', async () => { + const x = [[[2, 4]], [[9, 23]]]; // 2x1x2 + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const result = + tf.batchNorm3d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('batchnorm3D, x,mean,var,offset,scale are all 3D', async () => { + const shape: [number, number, number] = [2, 1, 2]; + const xT = tf.tensor3d([2, 4, 9, 23], shape); + const meanT = tf.tensor3d([1, 2, 3, 4], shape); + const varianceT = tf.tensor3d([2, 3, 4, 5], shape); + const offsetT = tf.tensor3d([3, 4, 5, 6], shape); + const scaleT = tf.tensor3d([4, 5, 6, 7], shape); + + const varianceEpsilon = .001; + + const result = + tf.batchNorm3d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const offset = await offsetT.buffer(); + const scale = await scaleT.buffer(); + expectArraysClose(await result.data(), [ + offset.get(0, 0, 0) + + (x.get(0, 0, 0) - mean.get(0, 0, 0)) * scale.get(0, 0, 0) / + Math.sqrt(variance.get(0, 0, 0) + varianceEpsilon), + offset.get(0, 0, 1) + + (x.get(0, 0, 1) - mean.get(0, 0, 1)) * scale.get(0, 0, 1) / + Math.sqrt(variance.get(0, 0, 1) + varianceEpsilon), + offset.get(1, 0, 0) + + (x.get(1, 0, 0) - mean.get(1, 0, 0)) * scale.get(1, 0, 0) / + Math.sqrt(variance.get(1, 0, 0) + varianceEpsilon), + offset.get(1, 0, 1) + + (x.get(1, 0, 1) - mean.get(1, 0, 1)) * scale.get(1, 0, 1) / + Math.sqrt(variance.get(1, 0, 1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm3D gradients, 2x1x2', async () => { + const x = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const offset = tf.tensor1d([3, 4]); + const scale = tf.tensor1d([2, 5]); + + const varianceEpsilon = .001; + + const dy = tf.tensor3d([1, 1, 1, 1], [2, 1, 2]); + const gradX = tf.grad( + (x: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); + expect(gradX.shape).toEqual([2, 1, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor1D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [-2.828, -5.773]); + expect(gradMean.shape).toEqual([2]); + const gradVariance = tf.grad( + (variance: tf.Tensor1D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); + expect(gradVariance.shape).toEqual([2]); + const gradOffset = tf.grad( + (offset: tf.Tensor1D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), [2, 2]); + expect(gradOffset.shape).toEqual([2]); + const gradScale = tf.grad( + (scale: tf.Tensor1D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [6.362, 13.277]); + expect(gradScale.shape).toEqual([2]); + }); + + it('batchnorm3D gradients, same shapes in x, mean and variance', async () => { + const x = tf.tensor3d([10, 20, 30, 40], [2, 1, 2]); + const mean = tf.tensor3d([0, 5, 10, 15], [2, 1, 2]); + const variance = tf.tensor3d([2, 4, 6, 8], [2, 1, 2]); + const scale = tf.tensor3d([2, 5, 2, 5], [2, 1, 2]); + const offset = tf.tensor3d([0, 0, 0, 0], [2, 1, 2]); + + const varianceEpsilon = .001; + + const dy = tf.tensor3d([1, 1, 1, 1], [2, 1, 2]); + const gradX = tf.grad( + (x: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [1.414, 2.500, 0.816, 1.768]); + expect(gradX.shape).toEqual([2, 1, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [-1.414, -2.500, -0.816, -1.768]); + expect(gradMean.shape).toEqual([2, 1, 2]); + const gradVariance = tf.grad( + (variance: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose( + await gradVariance.data(), [-3.533, -4.686, -1.360, -2.762]); + expect(gradVariance.shape).toEqual([2, 1, 2]); + const gradOffset = tf.grad( + (offset: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), [1, 1, 1, 1]); + expect(gradOffset.shape).toEqual([2, 1, 2]); + const gradScale = tf.grad( + (scale: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [7.069, 7.499, 8.164, 8.838]); + expect(gradScale.shape).toEqual([2, 1, 2]); + }); + + it('batchnorm matches tensorflow, 2x3x3', async () => { + const x = tf.tensor3d( + [ + 0.49955603, 0.04158615, -1.09440524, 2.03854165, -0.61578344, + 2.87533573, 1.18105987, 0.807462, 1.87888837, 2.26563962, -0.37040935, + 1.35848753, -0.75347094, 0.15683117, 0.91925946, 0.34121279, + 0.92717143, 1.89683965 + ], + [2, 3, 3]); + const mean = tf.tensor1d([0.39745062, -0.48062894, 0.4847822]); + const variance = tf.tensor1d([0.32375343, 0.67117643, 1.08334653]); + const offset = tf.tensor1d([0.69398749, -1.29056387, 0.9429723]); + const scale = tf.tensor1d([-0.5607271, 0.9878457, 0.25181573]); + const varianceEpsilon = .001; + + const result = + tf.batchNorm3d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + 0.59352049, -0.66135202, 0.5610874, -0.92077015, -1.45341019, 1.52106473, + -0.07704776, 0.26144429, 1.28010017, -1.14422404, -1.15776136, 1.15425493, + 1.82644104, -0.52249442, 1.04803919, 0.74932291, 0.40568101, 1.2844412 + ]); + }); +}); diff --git a/tfjs-core/src/ops/batchnorm4d.ts b/tfjs-core/src/ops/batchnorm4d.ts new file mode 100644 index 00000000000..5304cc7386d --- /dev/null +++ b/tfjs-core/src/ops/batchnorm4d.ts @@ -0,0 +1,93 @@ +/** + * @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 {Tensor1D, Tensor4D} from '../tensor'; +import {convertToTensor} from '../tensor_util_env'; +import {TensorLike} from '../types'; +import * as util from '../util'; + +import {op} from './operation'; +import {batchNorm} from './batchnorm'; +import {warnDeprecation} from './batchnorm_util'; +/** + * Batch normalization, strictly for 4D. For the more relaxed version, see + * `tf.batchNorm`. + * + * @param x The input Tensor. + * @param mean A mean Tensor. + * @param variance A variance Tensor. + * @param offset An offset Tensor. + * @param scale A scale Tensor. + * @param varianceEpsilon A small float number to avoid dividing by 0. + */ +function batchNorm4d_( + x: Tensor4D | TensorLike, mean: Tensor4D | Tensor1D | TensorLike, + variance: Tensor4D | Tensor1D | TensorLike, + offset?: Tensor4D | Tensor1D | TensorLike, scale?: Tensor4D | Tensor1D | TensorLike, + varianceEpsilon?: number): Tensor4D { + const $x = convertToTensor(x, 'x', 'batchNorm'); + const $mean = convertToTensor(mean, 'mean', 'batchNorm'); + const $variance = convertToTensor(variance, 'variance', 'batchNorm'); + let $scale: Tensor4D | Tensor1D; + if (scale != null) { + $scale = convertToTensor(scale, 'scale', 'batchNorm'); + } + let $offset: Tensor4D | Tensor1D; + if (offset != null) { + $offset = convertToTensor(offset, 'offset', 'batchNorm'); + } + util.assert( + $x.rank === 4, + () => `Error in batchNorm4D: x must be rank 4 but got rank ` + + `${$x.rank}.`); + util.assert( + $mean.rank === 4 || $mean.rank === 1, + () => `Error in batchNorm4D: mean must be rank 4 or rank 1 but ` + + `got rank ${$mean.rank}.`); + util.assert( + $variance.rank === 4 || $variance.rank === 1, + () => `Error in batchNorm4D: variance must be rank 4 or rank 1 ` + + `but got rank ${$variance.rank}.`); + if ($scale != null) { + util.assert( + $scale.rank === 4 || $scale.rank === 1, + () => `Error in batchNorm4D: scale must be rank 4 or rank 1 ` + + `but got rank ${$scale.rank}.`); + } + if ($offset != null) { + util.assert( + $offset.rank === 4 || $offset.rank === 1, + () => `Error in batchNorm4D: offset must be rank 4 or rank 1 ` + + `but got rank ${$offset.rank}.`); + } + return batchNorm($x, $mean, $variance, $offset, $scale, varianceEpsilon); +} + +/** + * @deprecated Please use `tf.batchNorm4d` instead and note the positional + * argument change of scale, offset, and varianceEpsilon. + */ +function batchNormalization4d_( + x: Tensor4D | TensorLike, mean: Tensor4D | Tensor1D | TensorLike, + variance: Tensor4D | Tensor1D | TensorLike, varianceEpsilon = .001, + scale?: Tensor4D | Tensor1D | TensorLike, + offset?: Tensor4D | Tensor1D | TensorLike): Tensor4D { + warnDeprecation(); + return batchNorm4d_(x, mean, variance, offset, scale, varianceEpsilon); +} + +export const batchNormalization4d = op({batchNormalization4d_}); +export const batchNorm4d = op({batchNorm4d_}); diff --git a/tfjs-core/src/ops/batchnorm4d_test.ts b/tfjs-core/src/ops/batchnorm4d_test.ts new file mode 100644 index 00000000000..60464c5630e --- /dev/null +++ b/tfjs-core/src/ops/batchnorm4d_test.ts @@ -0,0 +1,233 @@ +/** + * @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 * as tf from '../index'; +import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; +import {expectArraysClose} from '../test_util'; + +describeWithFlags('batchNorm4D', ALL_ENVS, () => { + it('simple batchnorm4D, no offset or scale, 2x1x1x2', async () => { + const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const varianceEpsilon = .001; + + const result = tf.batchNorm4d( + xT, meanT, varianceT, undefined, undefined, varianceEpsilon); + + const x = await xT.array(); + const mean = await meanT.array(); + const variance = await varianceT.array(); + expectArraysClose(await result.data(), [ + (x[0][0][0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), + (x[0][0][0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon), + (x[1][0][0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), + (x[1][0][0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('simple batchnorm4D, no offset, 2x1x1x2', async () => { + const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const scaleT = tf.tensor1d([4, 5]); + const varianceEpsilon = .001; + + const result = tf.batchNorm4d( + xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const scale = await scaleT.buffer(); + + expectArraysClose(await result.data(), [ + (x.get(0, 0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm4D, no scale, 2x1x1x2', async () => { + const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = tf.batchNorm4d( + xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const offset = await offsetT.buffer(); + + expectArraysClose(await result.data(), [ + offset.get(0) + + (x.get(0, 0, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(0, 0, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon), + offset.get(0) + + (x.get(1, 0, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(1, 0, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm4D, 2x1x1x2', async () => { + const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([3, 4]); + const scaleT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = + tf.batchNorm4d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const scale = await scaleT.buffer(); + const offset = await offsetT.buffer(); + + expectArraysClose(await result.data(), [ + offset.get(0) + + (x.get(0, 0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(0, 0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + offset.get(0) + + (x.get(1, 0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(1, 0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('accepts a tensor-like object', async () => { + const x = [[[[2, 4]]], [[[9, 23]]]]; // 2x1x1x2 + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const result = + tf.batchNorm4d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0][0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][0][0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0][0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][0][0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('simple batchnorm4D gradients, 2x1x1x2', async () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const offset = tf.tensor1d([3, 4]); + const scale = tf.tensor1d([2, 5]); + + const varianceEpsilon = .001; + + const dy = tf.tensor4d([-1, -1, -1, -1], [2, 1, 1, 2]); + const gradX = tf.grad( + (x: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [-1.414, -2.887, -1.414, -2.887]); + expect(gradX.shape).toEqual([2, 1, 1, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor1D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [2.828, 5.773]); + expect(gradMean.shape).toEqual([2]); + const gradVariance = tf.grad( + (variance: tf.Tensor1D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose(await gradVariance.data(), [3.180, 11.060]); + expect(gradVariance.shape).toEqual([2]); + const gradOffset = tf.grad( + (offset: tf.Tensor1D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), await dy.sum([0, 1, 2]).data()); + expect(gradOffset.shape).toEqual([2]); + const gradScale = tf.grad( + (scale: tf.Tensor1D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [-6.362, -13.277]); + expect(gradScale.shape).toEqual([2]); + }); + + it('batchnorm4D gradients, same shapes in x, mean and variance', async () => { + const x = tf.tensor4d([10, 20, 30, 40], [2, 1, 1, 2]); + const mean = tf.tensor4d([0, 5, 10, 15], [2, 1, 1, 2]); + const variance = tf.tensor4d([2, 4, 6, 8], [2, 1, 1, 2]); + const scale = tf.tensor4d([2, 5, 2, 5], [2, 1, 1, 2]); + const offset = tf.tensor4d([0, 0, 0, 0], [2, 1, 1, 2]); + + const varianceEpsilon = .001; + + const dy = tf.tensor4d([-1, -1, -1, -1], [2, 1, 1, 2]); + const gradX = tf.grad( + (x: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [-1.414, -2.500, -0.816, -1.768]); + expect(gradX.shape).toEqual([2, 1, 1, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [1.414, 2.500, 0.816, 1.768]); + expect(gradMean.shape).toEqual([2, 1, 1, 2]); + const gradVariance = tf.grad( + (variance: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose(await gradVariance.data(), [3.533, 4.686, 1.360, 2.762]); + expect(gradVariance.shape).toEqual([2, 1, 1, 2]); + const gradOffset = tf.grad( + (offset: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), await dy.data()); + expect(gradOffset.shape).toEqual([2, 1, 1, 2]); + const gradScale = tf.grad( + (scale: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [-7.069, -7.499, -8.164, -8.838]); + expect(gradScale.shape).toEqual([2, 1, 1, 2]); + }); +}); diff --git a/tfjs-core/src/ops/batchnorm_test.ts b/tfjs-core/src/ops/batchnorm_test.ts deleted file mode 100644 index 3b2e88b0f54..00000000000 --- a/tfjs-core/src/ops/batchnorm_test.ts +++ /dev/null @@ -1,950 +0,0 @@ -/** - * @license - * Copyright 2017 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('batchNorm4D', ALL_ENVS, () => { - it('simple batchnorm4D, no offset or scale, 2x1x1x2', async () => { - const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const varianceEpsilon = .001; - - const result = tf.batchNorm4d( - xT, meanT, varianceT, undefined, undefined, varianceEpsilon); - - const x = await xT.array(); - const mean = await meanT.array(); - const variance = await varianceT.array(); - expectArraysClose(await result.data(), [ - (x[0][0][0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), - (x[0][0][0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon), - (x[1][0][0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), - (x[1][0][0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('simple batchnorm4D, no offset, 2x1x1x2', async () => { - const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const scaleT = tf.tensor1d([4, 5]); - const varianceEpsilon = .001; - - const result = tf.batchNorm4d( - xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const scale = await scaleT.buffer(); - - expectArraysClose(await result.data(), [ - (x.get(0, 0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm4D, no scale, 2x1x1x2', async () => { - const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = tf.batchNorm4d( - xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const offset = await offsetT.buffer(); - - expectArraysClose(await result.data(), [ - offset.get(0) + - (x.get(0, 0, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(0, 0, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon), - offset.get(0) + - (x.get(1, 0, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(1, 0, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm4D, 2x1x1x2', async () => { - const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([3, 4]); - const scaleT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = - tf.batchNorm4d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const scale = await scaleT.buffer(); - const offset = await offsetT.buffer(); - - expectArraysClose(await result.data(), [ - offset.get(0) + - (x.get(0, 0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(0, 0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - offset.get(0) + - (x.get(1, 0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(1, 0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('accepts a tensor-like object', async () => { - const x = [[[[2, 4]]], [[[9, 23]]]]; // 2x1x1x2 - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const result = - tf.batchNorm4d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0][0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][0][0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0][0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][0][0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('simple batchnorm4D gradients, 2x1x1x2', async () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const offset = tf.tensor1d([3, 4]); - const scale = tf.tensor1d([2, 5]); - - const varianceEpsilon = .001; - - const dy = tf.tensor4d([-1, -1, -1, -1], [2, 1, 1, 2]); - const gradX = tf.grad( - (x: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [-1.414, -2.887, -1.414, -2.887]); - expect(gradX.shape).toEqual([2, 1, 1, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor1D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [2.828, 5.773]); - expect(gradMean.shape).toEqual([2]); - const gradVariance = tf.grad( - (variance: tf.Tensor1D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose(await gradVariance.data(), [3.180, 11.060]); - expect(gradVariance.shape).toEqual([2]); - const gradOffset = tf.grad( - (offset: tf.Tensor1D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), await dy.sum([0, 1, 2]).data()); - expect(gradOffset.shape).toEqual([2]); - const gradScale = tf.grad( - (scale: tf.Tensor1D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [-6.362, -13.277]); - expect(gradScale.shape).toEqual([2]); - }); - - it('batchnorm4D gradients, same shapes in x, mean and variance', async () => { - const x = tf.tensor4d([10, 20, 30, 40], [2, 1, 1, 2]); - const mean = tf.tensor4d([0, 5, 10, 15], [2, 1, 1, 2]); - const variance = tf.tensor4d([2, 4, 6, 8], [2, 1, 1, 2]); - const scale = tf.tensor4d([2, 5, 2, 5], [2, 1, 1, 2]); - const offset = tf.tensor4d([0, 0, 0, 0], [2, 1, 1, 2]); - - const varianceEpsilon = .001; - - const dy = tf.tensor4d([-1, -1, -1, -1], [2, 1, 1, 2]); - const gradX = tf.grad( - (x: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [-1.414, -2.500, -0.816, -1.768]); - expect(gradX.shape).toEqual([2, 1, 1, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [1.414, 2.500, 0.816, 1.768]); - expect(gradMean.shape).toEqual([2, 1, 1, 2]); - const gradVariance = tf.grad( - (variance: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose(await gradVariance.data(), [3.533, 4.686, 1.360, 2.762]); - expect(gradVariance.shape).toEqual([2, 1, 1, 2]); - const gradOffset = tf.grad( - (offset: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), await dy.data()); - expect(gradOffset.shape).toEqual([2, 1, 1, 2]); - const gradScale = tf.grad( - (scale: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [-7.069, -7.499, -8.164, -8.838]); - expect(gradScale.shape).toEqual([2, 1, 1, 2]); - }); -}); - -describeWithFlags('batchNorm3D', ALL_ENVS, () => { - it('simple batchnorm3D, no offset or scale, 2x1x2', async () => { - const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const varianceEpsilon = .001; - - const result = tf.batchNorm3d( - xT, meanT, varianceT, undefined, undefined, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - expectArraysClose(await result.data(), [ - (x.get(0, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm3D, no offset, 2x1x2', async () => { - const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const scaleT = tf.tensor1d([4, 5]); - const varianceEpsilon = .001; - - const result = tf.batchNorm3d( - xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const scale = await scaleT.buffer(); - expectArraysClose(await result.data(), [ - (x.get(0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm3D, no scale, 2x1x2', async () => { - const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = tf.batchNorm3d( - xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const offset = await offsetT.buffer(); - expectArraysClose(await result.data(), [ - offset.get(0) + - (x.get(0, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(0, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon), - offset.get(0) + - (x.get(1, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(1, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm3D, 2x1x2', async () => { - const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([3, 4]); - const scaleT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = - tf.batchNorm3d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const offset = await offsetT.buffer(); - const scale = await scaleT.buffer(); - - expectArraysClose(await result.data(), [ - offset.get(0) + - (x.get(0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - offset.get(0) + - (x.get(1, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(1, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('accepts a tensor-like object', async () => { - const x = [[[2, 4]], [[9, 23]]]; // 2x1x2 - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const result = - tf.batchNorm3d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('batchnorm3D, x,mean,var,offset,scale are all 3D', async () => { - const shape: [number, number, number] = [2, 1, 2]; - const xT = tf.tensor3d([2, 4, 9, 23], shape); - const meanT = tf.tensor3d([1, 2, 3, 4], shape); - const varianceT = tf.tensor3d([2, 3, 4, 5], shape); - const offsetT = tf.tensor3d([3, 4, 5, 6], shape); - const scaleT = tf.tensor3d([4, 5, 6, 7], shape); - - const varianceEpsilon = .001; - - const result = - tf.batchNorm3d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const offset = await offsetT.buffer(); - const scale = await scaleT.buffer(); - expectArraysClose(await result.data(), [ - offset.get(0, 0, 0) + - (x.get(0, 0, 0) - mean.get(0, 0, 0)) * scale.get(0, 0, 0) / - Math.sqrt(variance.get(0, 0, 0) + varianceEpsilon), - offset.get(0, 0, 1) + - (x.get(0, 0, 1) - mean.get(0, 0, 1)) * scale.get(0, 0, 1) / - Math.sqrt(variance.get(0, 0, 1) + varianceEpsilon), - offset.get(1, 0, 0) + - (x.get(1, 0, 0) - mean.get(1, 0, 0)) * scale.get(1, 0, 0) / - Math.sqrt(variance.get(1, 0, 0) + varianceEpsilon), - offset.get(1, 0, 1) + - (x.get(1, 0, 1) - mean.get(1, 0, 1)) * scale.get(1, 0, 1) / - Math.sqrt(variance.get(1, 0, 1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm3D gradients, 2x1x2', async () => { - const x = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const offset = tf.tensor1d([3, 4]); - const scale = tf.tensor1d([2, 5]); - - const varianceEpsilon = .001; - - const dy = tf.tensor3d([1, 1, 1, 1], [2, 1, 2]); - const gradX = tf.grad( - (x: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); - expect(gradX.shape).toEqual([2, 1, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor1D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [-2.828, -5.773]); - expect(gradMean.shape).toEqual([2]); - const gradVariance = tf.grad( - (variance: tf.Tensor1D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); - expect(gradVariance.shape).toEqual([2]); - const gradOffset = tf.grad( - (offset: tf.Tensor1D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), [2, 2]); - expect(gradOffset.shape).toEqual([2]); - const gradScale = tf.grad( - (scale: tf.Tensor1D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [6.362, 13.277]); - expect(gradScale.shape).toEqual([2]); - }); - - it('batchnorm3D gradients, same shapes in x, mean and variance', async () => { - const x = tf.tensor3d([10, 20, 30, 40], [2, 1, 2]); - const mean = tf.tensor3d([0, 5, 10, 15], [2, 1, 2]); - const variance = tf.tensor3d([2, 4, 6, 8], [2, 1, 2]); - const scale = tf.tensor3d([2, 5, 2, 5], [2, 1, 2]); - const offset = tf.tensor3d([0, 0, 0, 0], [2, 1, 2]); - - const varianceEpsilon = .001; - - const dy = tf.tensor3d([1, 1, 1, 1], [2, 1, 2]); - const gradX = tf.grad( - (x: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [1.414, 2.500, 0.816, 1.768]); - expect(gradX.shape).toEqual([2, 1, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [-1.414, -2.500, -0.816, -1.768]); - expect(gradMean.shape).toEqual([2, 1, 2]); - const gradVariance = tf.grad( - (variance: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose( - await gradVariance.data(), [-3.533, -4.686, -1.360, -2.762]); - expect(gradVariance.shape).toEqual([2, 1, 2]); - const gradOffset = tf.grad( - (offset: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), [1, 1, 1, 1]); - expect(gradOffset.shape).toEqual([2, 1, 2]); - const gradScale = tf.grad( - (scale: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [7.069, 7.499, 8.164, 8.838]); - expect(gradScale.shape).toEqual([2, 1, 2]); - }); - - it('batchnorm matches tensorflow, 2x3x3', async () => { - const x = tf.tensor3d( - [ - 0.49955603, 0.04158615, -1.09440524, 2.03854165, -0.61578344, - 2.87533573, 1.18105987, 0.807462, 1.87888837, 2.26563962, -0.37040935, - 1.35848753, -0.75347094, 0.15683117, 0.91925946, 0.34121279, - 0.92717143, 1.89683965 - ], - [2, 3, 3]); - const mean = tf.tensor1d([0.39745062, -0.48062894, 0.4847822]); - const variance = tf.tensor1d([0.32375343, 0.67117643, 1.08334653]); - const offset = tf.tensor1d([0.69398749, -1.29056387, 0.9429723]); - const scale = tf.tensor1d([-0.5607271, 0.9878457, 0.25181573]); - const varianceEpsilon = .001; - - const result = - tf.batchNorm3d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - 0.59352049, -0.66135202, 0.5610874, -0.92077015, -1.45341019, 1.52106473, - -0.07704776, 0.26144429, 1.28010017, -1.14422404, -1.15776136, 1.15425493, - 1.82644104, -0.52249442, 1.04803919, 0.74932291, 0.40568101, 1.2844412 - ]); - }); -}); - -describeWithFlags('batchNorm2D', ALL_ENVS, () => { - it('simple batchnorm2D, no offset or scale, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const varianceEpsilon = .001; - - const result = tf.batchNorm2d( - xT, meanT, varianceT, undefined, undefined, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - expectArraysClose(await result.data(), [ - (x.get(0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - it('simple batchnorm2D, no offset, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const scaleT = tf.tensor1d([4, 5]); - const varianceEpsilon = .001; - - const result = tf.batchNorm2d( - xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const scale = await scaleT.buffer(); - expectArraysClose(await result.data(), [ - (x.get(0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm2D, no scale, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = tf.batchNorm2d( - xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); - - const offset = await offsetT.array(); - const mean = await meanT.array(); - const variance = await varianceT.array(); - const x = await xT.array(); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('simple batchnorm2D, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([3, 4]); - const scaleT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = - tf.batchNorm2d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); - - const offset = await offsetT.array(); - const mean = await meanT.array(); - const variance = await varianceT.array(); - const scale = await scaleT.array(); - const x = await xT.array(); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('simple batchnorm2D gradients, 2x2', async () => { - const x = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const offset = tf.tensor1d([3, 4]); - const scale = tf.tensor1d([2, 5]); - - const varianceEpsilon = .001; - - const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); - const [gradX, gradMean, gradVariance, gradOffset, gradScale] = tf.grads( - (x: tf.Tensor2D, mean: tf.Tensor1D, variance: tf.Tensor1D, - offset: tf.Tensor1D, scale: tf.Tensor1D) => - tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon))( - [x, mean, variance, offset, scale], dy); - expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); - expect(gradX.shape).toEqual([2, 2]); - expectArraysClose(await gradMean.data(), [-2.828, -5.773]); - expect(gradMean.shape).toEqual([2]); - expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); - expect(gradVariance.shape).toEqual([2]); - expectArraysClose(await gradOffset.data(), [2, 2]); - expect(gradOffset.shape).toEqual([2]); - expectArraysClose(await gradScale.data(), [6.362, 13.277]); - expect(gradScale.shape).toEqual([2]); - }); - - it('gradient with clones batchnorm2D', async () => { - const x = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const offset = tf.tensor1d([3, 4]); - const scale = tf.tensor1d([2, 5]); - - const varianceEpsilon = .001; - - const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); - const [gradX, gradMean, gradVariance, gradOffset, gradScale] = tf.grads( - (x: tf.Tensor2D, mean: tf.Tensor1D, variance: tf.Tensor1D, - offset: tf.Tensor1D, scale: tf.Tensor1D) => - tf.batchNorm2d( - x.clone(), mean.clone(), variance.clone(), offset.clone(), - scale.clone(), varianceEpsilon) - .clone())([x, mean, variance, offset, scale], dy); - expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); - expect(gradX.shape).toEqual([2, 2]); - expectArraysClose(await gradMean.data(), [-2.828, -5.773]); - expect(gradMean.shape).toEqual([2]); - expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); - expect(gradVariance.shape).toEqual([2]); - expectArraysClose(await gradOffset.data(), [2, 2]); - expect(gradOffset.shape).toEqual([2]); - expectArraysClose(await gradScale.data(), [6.362, 13.277]); - expect(gradScale.shape).toEqual([2]); - }); - - it('batchnorm2D gradients, same shapes in x, mean and variance', async () => { - const x = tf.tensor2d([10, 20, 30, 40], [2, 2]); - const mean = tf.tensor2d([0, 5, 10, 15], [2, 2]); - const variance = tf.tensor2d([2, 4, 6, 8], [2, 2]); - const scale = tf.tensor2d([2, 5, 2, 5], [2, 2]); - const offset = tf.tensor2d([0, 0, 0, 0], [2, 2]); - - const varianceEpsilon = .001; - - const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); - const gradX = tf.grad( - (x: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [1.414, 2.500, 0.816, 1.768]); - expect(gradX.shape).toEqual([2, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [-1.414, -2.500, -0.816, -1.768]); - expect(gradMean.shape).toEqual([2, 2]); - const gradVariance = tf.grad( - (variance: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose( - await gradVariance.data(), [-3.533, -4.686, -1.360, -2.762]); - expect(gradVariance.shape).toEqual([2, 2]); - const gradOffset = tf.grad( - (offset: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), [1, 1, 1, 1]); - expect(gradOffset.shape).toEqual([2, 2]); - const gradScale = tf.grad( - (scale: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [7.069, 7.499, 8.164, 8.838]); - expect(gradScale.shape).toEqual([2, 2]); - }); - - it('gradient with clones', () => { - const x = tf.zeros([2, 2]); - const mean = tf.zeros([2, 2]); - const variance = tf.zeros([2, 2]); - const scale = tf.zeros([2, 2]); - const offset = tf.zeros([2, 2]); - - const varianceEpsilon = .001; - - const gradF = tf.grads( - (x: tf.Tensor2D, mean: tf.Tensor2D, variance: tf.Tensor2D, - offset: tf.Tensor2D, scale: tf.Tensor2D) => - tf.batchNorm2d( - x.clone(), mean.clone(), variance.clone(), offset.clone(), - scale.clone(), varianceEpsilon) - .clone()); - const [gradX, gradMean, gradVariance, gradOffset, gradScale] = - gradF([x, mean, variance, offset, scale]); - expect(gradX.shape).toEqual(x.shape); - expect(gradMean.shape).toEqual(mean.shape); - expect(gradVariance.shape).toEqual(variance.shape); - expect(gradOffset.shape).toEqual(offset.shape); - expect(gradScale.shape).toEqual(scale.shape); - }); - - it('batchnorm2D matches tensorflow, 3x3', async () => { - const x = tf.tensor2d( - [ - 0.3136892, 0.92389025, 0.594782, 0.05021042, 0.67545404, 0.93910035, - 0.13277993, 0.96474269, 0.88608916 - ], - [3, 3]); - const mean = tf.tensor1d([0.19526312, 0.74857256, 0.45166398]); - const variance = tf.tensor1d([0.22963001, 0.61521992, 0.46623685]); - const offset = tf.tensor1d([0.43098484, 0.77712237, 0.47916298]); - const scale = tf.tensor1d([0.62186907, 0.85673736, 0.19201061]); - const varianceEpsilon = .001; - - const result = - tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - 0.58433646, 0.96846228, 0.51936529, 0.24315402, 0.69732157, 0.61608542, - 0.35007446, 1.01304821, 0.60119441 - ]); - }); - - it('throws when passed x as a non-tensor', () => { - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - - expect(() => tf.batchNorm({} as tf.Tensor, mean, variance)) - .toThrowError(/Argument 'x' passed to 'batchNorm' must be a Tensor/); - }); - it('throws when passed mean as a non-tensor', () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const variance = tf.tensor1d([2, 3]); - - expect(() => tf.batchNorm(x, {} as tf.Tensor, variance)) - .toThrowError(/Argument 'mean' passed to 'batchNorm' must be a Tensor/); - }); - it('throws when passed variance as a non-tensor', () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const mean = tf.tensor1d([1, 2]); - - const e = /Argument 'variance' passed to 'batchNorm' must be a Tensor/; - expect(() => tf.batchNorm(x, mean, {} as tf.Tensor)).toThrowError(e); - }); - it('throws when passed scale as a non-tensor', () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const epsilon = .001; - - expect(() => tf.batchNorm(x, mean, variance, epsilon, {} as tf.Tensor)) - .toThrowError( - /Argument 'scale' passed to 'batchNorm' must be a Tensor/); - }); - it('throws when passed offset as a non-tensor', () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const epsilon = .001; - const scale = tf.tensor1d([0.62186907, 0.85673736, 0.19201061]); - - const e = /Argument 'offset' passed to 'batchNorm' must be a Tensor/; - expect( - () => tf.batchNorm(x, mean, variance, {} as tf.Tensor, scale, epsilon)) - .toThrowError(e); - }); - - it('accepts a tensor-like object', async () => { - const x = [[2, 4], [9, 23]]; - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const result = - tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('throws error when x is a string tensor', () => { - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const f = () => tf.batchNorm2d( - [['a', 'b'], ['c', 'd']], mean, variance, offset, scale, - varianceEpsilon); - expect(f).toThrowError( - /Argument 'x' passed to 'batchNorm' must be numeric/); - }); - - it('throws error when mean is a string tensor', () => { - const x = [[2, 4], [9, 23]]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const f = () => - tf.batchNorm2d(x, ['a', 'b'], variance, offset, scale, varianceEpsilon); - expect(f).toThrowError( - /Argument 'mean' passed to 'batchNorm' must be numeric/); - }); - - it('throws error when variance is a string tensor', () => { - const x = [[2, 4], [9, 23]]; - const mean = [1, 2]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const f = () => - tf.batchNorm2d(x, mean, ['a', 'b'], offset, scale, varianceEpsilon); - expect(f).toThrowError(/'variance' passed to 'batchNorm' must be numeric/); - }); - - it('throws error when scale is a string tensor', () => { - const x = [[2, 4], [9, 23]]; - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - - const varianceEpsilon = .001; - - const f = () => - tf.batchNorm2d(x, mean, variance, offset, ['a', 'b'], varianceEpsilon); - expect(f).toThrowError(/'scale' passed to 'batchNorm' must be numeric/); - }); - - it('throws error when offset is a string tensor', () => { - const x = [[2, 4], [9, 23]]; - const mean = [1, 2]; - const variance = [2, 3]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const f = () => - tf.batchNorm2d(x, mean, variance, ['a', 'b'], scale, varianceEpsilon); - expect(f).toThrowError(/'offset' passed to 'batchNorm' must be numeric/); - }); -}); - -describeWithFlags('deprecated batchNormalization', ALL_ENVS, () => { - beforeAll(() => { - // Silence deprecation warnings. - spyOn(console, 'warn'); - }); - - it('simple batchnorm2D, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([3, 4]); - const scaleT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = tf.batchNormalization( - xT, meanT, varianceT, varianceEpsilon, scaleT, offsetT); - - const offset = await offsetT.array(); - const mean = await meanT.array(); - const variance = await varianceT.array(); - const scale = await scaleT.array(); - const x = await xT.array(); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - - const result2 = tf.batchNormalization2d( - xT, meanT, varianceT, varianceEpsilon, scaleT, offsetT); - expectArraysClose(await result.data(), await result2.data()); - }); -}); diff --git a/tfjs-core/src/ops/batchnorm_util.ts b/tfjs-core/src/ops/batchnorm_util.ts new file mode 100644 index 00000000000..b2b4e2f74c3 --- /dev/null +++ b/tfjs-core/src/ops/batchnorm_util.ts @@ -0,0 +1,24 @@ +/** + * @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 {deprecationWarn} from '../globals'; + +export function warnDeprecation() { + deprecationWarn( + 'tf.batchNormalization() is going away. ' + + 'Use tf.batchNorm() instead, and note the positional argument change ' + + 'of scale, offset, and varianceEpsilon'); +} diff --git a/tfjs-core/src/ops/ops.ts b/tfjs-core/src/ops/ops.ts index 404e85097c4..0f45e6e42bf 100644 --- a/tfjs-core/src/ops/ops.ts +++ b/tfjs-core/src/ops/ops.ts @@ -16,6 +16,10 @@ */ // Modularized ops. +export {batchNormalization, batchNorm} from './batchnorm'; +export {batchNormalization2d, batchNorm2d} from './batchnorm2d'; +export {batchNormalization3d, batchNorm3d} from './batchnorm3d'; +export {batchNormalization4d, batchNorm4d} from './batchnorm4d'; export {broadcastTo} from './broadcast_to'; export {clone} from './clone'; export {div} from './div'; @@ -37,7 +41,6 @@ export {squaredDifference} from './squared_difference'; export {tile} from './tile'; export {truncatedNormal} from './truncated_normal'; -export * from './batchnorm'; export * from './boolean_mask'; export * from './complex_ops'; export * from './concat_split'; diff --git a/tfjs-core/src/public/chained_ops/batchnorm.ts b/tfjs-core/src/public/chained_ops/batchnorm.ts new file mode 100644 index 00000000000..0e26af18d97 --- /dev/null +++ b/tfjs-core/src/public/chained_ops/batchnorm.ts @@ -0,0 +1,40 @@ +/** + * @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 {batchNorm} from '../../ops/batchnorm'; +import {Tensor, Tensor1D} from '../../tensor'; +import {Rank, TensorLike} from '../../types'; + +declare module '../../tensor' { + interface Tensor { + batchNorm( + mean: Tensor | Tensor1D | TensorLike, + variance: Tensor | Tensor1D | TensorLike, + offset?: Tensor | Tensor1D | TensorLike, + scale?: Tensor | Tensor1D | TensorLike, + varianceEpsilon?: number + ): Tensor; + } +} + +Tensor.prototype.batchNorm = function ( + this: Tensor | TensorLike, mean: Tensor | Tensor1D | TensorLike, + variance: Tensor | Tensor1D | TensorLike, + offset?: Tensor | Tensor1D | TensorLike, + scale?: Tensor | Tensor1D | TensorLike, + varianceEpsilon?: number): Tensor { + return batchNorm(this, mean, variance, offset, scale, varianceEpsilon); +} 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 d398ba1944d..ca8aa8194db 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 @@ -23,3 +23,4 @@ import './tile'; import './one_hot'; import './transpose'; import './pad'; +import './batchnorm'; 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 92aa87067c0..8f1630653d2 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 @@ -25,7 +25,7 @@ import {ALL_ENVS, describeWithFlags} from '../../jasmine_util'; const CHAINED_OPS = [ 'square', 'broadcastTo', 'tile', 'oneHot', 'div', 'divNoNan', 'transpose', - 'pad' + 'pad', 'batchNorm' ]; describeWithFlags('chained ops', ALL_ENVS, () => { @@ -34,7 +34,7 @@ describeWithFlags('chained ops', ALL_ENVS, () => { for (const opName of CHAINED_OPS) { //@ts-ignore expect(typeof tensor[opName]) - .toBe('function', `${opName} chained op not found`); + .toBe('function', `${opName} chained op not found`); } }); }); diff --git a/tfjs-core/src/tensor.ts b/tfjs-core/src/tensor.ts index c7036e2f3a5..83345f02555 100644 --- a/tfjs-core/src/tensor.ts +++ b/tfjs-core/src/tensor.ts @@ -195,12 +195,6 @@ export interface OpHandler { concat(tensors: Array, axis: number): T; stack(tensors: Array, axis: number): Tensor; unstack(value: T, axis: number): Tensor[]; - batchNorm( - x: Tensor, mean: Tensor|Tensor1D|TensorLike, - variance: Tensor|Tensor1D|TensorLike, - offset?: Tensor|Tensor1D|TensorLike, - scale?: Tensor|Tensor1D|TensorLike, - varianceEpsilon?: number): Tensor; all(x: Tensor, axis: number|number[], keepDims: boolean): T; any(x: Tensor, axis: number|number[], keepDims: boolean): T; logSumExp( @@ -833,17 +827,6 @@ export class Tensor { return this.batchNorm(mean, variance, offset, scale, varianceEpsilon); } - batchNorm( - mean: Tensor|Tensor1D|TensorLike, - variance: Tensor|Tensor1D|TensorLike, - offset?: Tensor|Tensor1D|TensorLike, - scale?: Tensor|Tensor1D|TensorLike, - varianceEpsilon = .001, - ): Tensor { - this.throwIfDisposed(); - return opHandler.batchNorm( - this, mean, variance, offset, scale, varianceEpsilon); - } // Reduction ops. all(axis: number|number[] = null, keepDims = false): T { this.throwIfDisposed(); diff --git a/tfjs-core/src/tests.ts b/tfjs-core/src/tests.ts index ae9143836b5..7ddf226263f 100644 --- a/tfjs-core/src/tests.ts +++ b/tfjs-core/src/tests.ts @@ -40,7 +40,9 @@ import './kernel_registry_test'; import './ops/arithmetic_test'; import './ops/array_ops_test'; import './ops/axis_util_test'; -import './ops/batchnorm_test'; +import './ops/batchnorm2d_test'; +import './ops/batchnorm3d_test'; +import './ops/batchnorm4d_test'; import './ops/binary_ops_test'; import './ops/boolean_mask_test'; import './ops/broadcast_to_test'; From 3f44a44062d65389b3d757a1186e7dfdbde98463 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 30 Mar 2020 14:09:09 -0700 Subject: [PATCH 2/8] Fix lint. --- .vscode/settings.json | 3 +- tfjs-core/src/ops/batchnorm.ts | 78 +- tfjs-core/src/ops/batchnorm2d.ts | 3 +- tfjs-core/src/ops/batchnorm2d_test.ts | 458 --------- tfjs-core/src/ops/batchnorm3d.ts | 3 +- tfjs-core/src/ops/batchnorm3d_test.ts | 296 ------ tfjs-core/src/ops/batchnorm4d.ts | 3 +- tfjs-core/src/ops/batchnorm4d_test.ts | 233 ----- tfjs-core/src/ops/batchnorm_test.ts | 950 ++++++++++++++++++ tfjs-core/src/public/chained_ops/batchnorm.ts | 2 +- tfjs-core/src/tests.ts | 4 +- 11 files changed, 999 insertions(+), 1034 deletions(-) delete mode 100644 tfjs-core/src/ops/batchnorm2d_test.ts delete mode 100644 tfjs-core/src/ops/batchnorm3d_test.ts delete mode 100644 tfjs-core/src/ops/batchnorm4d_test.ts create mode 100644 tfjs-core/src/ops/batchnorm_test.ts diff --git a/.vscode/settings.json b/.vscode/settings.json index 2a84b4c07db..997d33d2d42 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -18,7 +18,8 @@ "editor.tabSize": 2, "editor.insertSpaces": true, "[typescript]": { - "editor.formatOnSave": true + "editor.formatOnSave": true, + "editor.defaultFormatter": "xaver.clang-format" }, "[javascript]": { "editor.formatOnSave": true diff --git a/tfjs-core/src/ops/batchnorm.ts b/tfjs-core/src/ops/batchnorm.ts index 317da3a579c..ecbb7b42323 100644 --- a/tfjs-core/src/ops/batchnorm.ts +++ b/tfjs-core/src/ops/batchnorm.ts @@ -21,22 +21,22 @@ import {convertToTensor} from '../tensor_util_env'; import {Rank, ShapeMap, TensorLike} from '../types'; import * as util from '../util'; +import {warnDeprecation} from './batchnorm_util'; import {getReductionAxes} from './broadcast_util'; import {op} from './operation'; import {scalar} from './tensor_ops'; import {tile} from './tile'; import {rsqrt} from './unary_ops'; -import {warnDeprecation} from './batchnorm_util'; /** * @deprecated Please use `tf.batchNorm` instead and note the positional * argument change of scale, offset, and varianceEpsilon. */ function batchNormalization_( - x: Tensor | TensorLike, mean: Tensor | Tensor1D | TensorLike, - variance: Tensor | Tensor1D | TensorLike, varianceEpsilon = .001, - scale?: Tensor | Tensor1D | TensorLike, - offset?: Tensor | Tensor1D | TensorLike): Tensor { + x: Tensor|TensorLike, mean: Tensor|Tensor1D|TensorLike, + variance: Tensor|Tensor1D|TensorLike, varianceEpsilon = .001, + scale?: Tensor|Tensor1D|TensorLike, + offset?: Tensor|Tensor1D|TensorLike): Tensor { warnDeprecation(); return batchNorm_(x, mean, variance, offset, scale, varianceEpsilon); } @@ -67,38 +67,38 @@ function batchNormalization_( */ /** @doc {heading: 'Operations', subheading: 'Normalization'} */ function batchNorm_( - x: Tensor | TensorLike, mean: Tensor | Tensor1D | TensorLike, - variance: Tensor | Tensor1D | TensorLike, - offset?: Tensor | Tensor1D | TensorLike, - scale?: Tensor | Tensor1D | TensorLike, - varianceEpsilon?: number): Tensor { + x: Tensor|TensorLike, mean: Tensor|Tensor1D|TensorLike, + variance: Tensor|Tensor1D|TensorLike, + offset?: Tensor|Tensor1D|TensorLike, + scale?: Tensor|Tensor1D|TensorLike, + varianceEpsilon?: number): Tensor { if (varianceEpsilon == null) { varianceEpsilon = 0.001; } const $x = convertToTensor(x, 'x', 'batchNorm'); const $mean = convertToTensor(mean, 'mean', 'batchNorm'); const $variance = convertToTensor(variance, 'variance', 'batchNorm'); - let $scale: Tensor | Tensor1D; + let $scale: Tensor|Tensor1D; if (scale != null) { $scale = convertToTensor(scale, 'scale', 'batchNorm'); } - let $offset: Tensor | Tensor1D; + let $offset: Tensor|Tensor1D; if (offset != null) { $offset = convertToTensor(offset, 'offset', 'batchNorm'); } util.assert( - $mean.rank === $variance.rank, - () => 'Batch normalization gradient requires mean and variance to have ' + - 'equal ranks.'); + $mean.rank === $variance.rank, + () => 'Batch normalization gradient requires mean and variance to have ' + + 'equal ranks.'); util.assert( - $offset == null || $mean.rank === $offset.rank, - () => 'Batch normalization gradient requires mean and offset to have ' + - 'equal ranks.'); + $offset == null || $mean.rank === $offset.rank, + () => 'Batch normalization gradient requires mean and offset to have ' + + 'equal ranks.'); util.assert( - $scale == null || $mean.rank === $scale.rank, - () => 'Batch normalization gradient requires mean and scale to have ' + - 'equal ranks.'); + $scale == null || $mean.rank === $scale.rank, + () => 'Batch normalization gradient requires mean and scale to have ' + + 'equal ranks.'); let x4D: Tensor4D; if ($x.rank === 0 || $x.rank === 1) { @@ -113,7 +113,7 @@ function batchNorm_( const der = (dy: Tensor, saved: Tensor[]) => { type Saved = [ - Tensor, Tensor | Tensor1D, Tensor | Tensor1D, Tensor | Tensor1D + Tensor, Tensor| Tensor1D, Tensor| Tensor1D, Tensor| Tensor1D ]; const [$x, $mean, $variance, $scale] = saved as Saved; const scaleValue = $scale == null ? scalar(1) : $scale; @@ -130,16 +130,16 @@ function batchNorm_( const dyTimesScaleValue = dy.mul(scaleValue); const oneOverSqrtVariance = rsqrt($variance.add(scalar(varianceEpsilon))); const minusHalfRCube = oneOverSqrtVariance.mul(oneOverSqrtVariance) - .mul(oneOverSqrtVariance) - .mul(scalar(-0.5)); + .mul(oneOverSqrtVariance) + .mul(scalar(-0.5)); const derX = () => { if ($mean.rank === 1) { return dy - .mul(tile( - oneOverSqrtVariance.as4D(1, 1, 1, $mean.shape[0]), tileShape)) - .mul(scaleValue) - .reshape($x.shape); + .mul(tile( + oneOverSqrtVariance.as4D(1, 1, 1, $mean.shape[0]), tileShape)) + .mul(scaleValue) + .reshape($x.shape); } else { return dy.mul(oneOverSqrtVariance).mul(scaleValue).reshape($x.shape); } @@ -185,20 +185,20 @@ function batchNorm_( const inputsToSave = [$x, $mean, $variance, $scale]; const res = ENGINE.runKernelFunc( - (backend, save) => { - const res = backend.batchNormalization( - x4D, batchnormReshape4D($mean), batchnormReshape4D($variance), - varianceEpsilon, batchnormReshape4D($scale), - batchnormReshape4D($offset)); - save([$x, $mean, $variance, $scale]); - return res; - }, - {x: $x, mean: $mean, variance: $variance, scale: $scale, offset: $offset}, - der, 'BatchNormalization', {varianceEpsilon}, inputsToSave); + (backend, save) => { + const res = backend.batchNormalization( + x4D, batchnormReshape4D($mean), batchnormReshape4D($variance), + varianceEpsilon, batchnormReshape4D($scale), + batchnormReshape4D($offset)); + save([$x, $mean, $variance, $scale]); + return res; + }, + {x: $x, mean: $mean, variance: $variance, scale: $scale, offset: $offset}, + der, 'BatchNormalization', {varianceEpsilon}, inputsToSave); return res.reshape($x.shape); } -function batchnormReshape4D(x: Tensor): Tensor4D | Tensor1D { +function batchnormReshape4D(x: Tensor): Tensor4D|Tensor1D { if (x == null) { return null; } diff --git a/tfjs-core/src/ops/batchnorm2d.ts b/tfjs-core/src/ops/batchnorm2d.ts index 7f8fe25ae84..09711430d01 100644 --- a/tfjs-core/src/ops/batchnorm2d.ts +++ b/tfjs-core/src/ops/batchnorm2d.ts @@ -37,7 +37,8 @@ import {warnDeprecation} from './batchnorm_util'; function batchNorm2d_( x: Tensor2D | TensorLike, mean: Tensor2D | Tensor1D | TensorLike, variance: Tensor2D | Tensor1D | TensorLike, - offset?: Tensor2D | Tensor1D | TensorLike, scale?: Tensor2D | Tensor1D | TensorLike, + offset?: Tensor2D | Tensor1D | TensorLike, + scale?: Tensor2D | Tensor1D | TensorLike, varianceEpsilon?: number): Tensor2D { const $x = convertToTensor(x, 'x', 'batchNorm'); const $mean = convertToTensor(mean, 'mean', 'batchNorm'); diff --git a/tfjs-core/src/ops/batchnorm2d_test.ts b/tfjs-core/src/ops/batchnorm2d_test.ts deleted file mode 100644 index 2eb33fe0d64..00000000000 --- a/tfjs-core/src/ops/batchnorm2d_test.ts +++ /dev/null @@ -1,458 +0,0 @@ -/** - * @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 * as tf from '../index'; -import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysClose} from '../test_util'; - -describeWithFlags('batchNorm2D', ALL_ENVS, () => { - it('simple batchnorm2D, no offset or scale, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const varianceEpsilon = .001; - - const result = tf.batchNorm2d( - xT, meanT, varianceT, undefined, undefined, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - expectArraysClose(await result.data(), [ - (x.get(0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - it('simple batchnorm2D, no offset, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const scaleT = tf.tensor1d([4, 5]); - const varianceEpsilon = .001; - - const result = tf.batchNorm2d( - xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const scale = await scaleT.buffer(); - expectArraysClose(await result.data(), [ - (x.get(0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm2D, no scale, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = tf.batchNorm2d( - xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); - - const offset = await offsetT.array(); - const mean = await meanT.array(); - const variance = await varianceT.array(); - const x = await xT.array(); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('simple batchnorm2D, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([3, 4]); - const scaleT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = - tf.batchNorm2d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); - - const offset = await offsetT.array(); - const mean = await meanT.array(); - const variance = await varianceT.array(); - const scale = await scaleT.array(); - const x = await xT.array(); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('simple batchnorm2D gradients, 2x2', async () => { - const x = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const offset = tf.tensor1d([3, 4]); - const scale = tf.tensor1d([2, 5]); - - const varianceEpsilon = .001; - - const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); - const [gradX, gradMean, gradVariance, gradOffset, gradScale] = tf.grads( - (x: tf.Tensor2D, mean: tf.Tensor1D, variance: tf.Tensor1D, - offset: tf.Tensor1D, scale: tf.Tensor1D) => - tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon))( - [x, mean, variance, offset, scale], dy); - expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); - expect(gradX.shape).toEqual([2, 2]); - expectArraysClose(await gradMean.data(), [-2.828, -5.773]); - expect(gradMean.shape).toEqual([2]); - expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); - expect(gradVariance.shape).toEqual([2]); - expectArraysClose(await gradOffset.data(), [2, 2]); - expect(gradOffset.shape).toEqual([2]); - expectArraysClose(await gradScale.data(), [6.362, 13.277]); - expect(gradScale.shape).toEqual([2]); - }); - - it('gradient with clones batchnorm2D', async () => { - const x = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const offset = tf.tensor1d([3, 4]); - const scale = tf.tensor1d([2, 5]); - - const varianceEpsilon = .001; - - const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); - const [gradX, gradMean, gradVariance, gradOffset, gradScale] = tf.grads( - (x: tf.Tensor2D, mean: tf.Tensor1D, variance: tf.Tensor1D, - offset: tf.Tensor1D, scale: tf.Tensor1D) => - tf.batchNorm2d( - x.clone(), mean.clone(), variance.clone(), offset.clone(), - scale.clone(), varianceEpsilon) - .clone())([x, mean, variance, offset, scale], dy); - expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); - expect(gradX.shape).toEqual([2, 2]); - expectArraysClose(await gradMean.data(), [-2.828, -5.773]); - expect(gradMean.shape).toEqual([2]); - expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); - expect(gradVariance.shape).toEqual([2]); - expectArraysClose(await gradOffset.data(), [2, 2]); - expect(gradOffset.shape).toEqual([2]); - expectArraysClose(await gradScale.data(), [6.362, 13.277]); - expect(gradScale.shape).toEqual([2]); - }); - - it('batchnorm2D gradients, same shapes in x, mean and variance', async () => { - const x = tf.tensor2d([10, 20, 30, 40], [2, 2]); - const mean = tf.tensor2d([0, 5, 10, 15], [2, 2]); - const variance = tf.tensor2d([2, 4, 6, 8], [2, 2]); - const scale = tf.tensor2d([2, 5, 2, 5], [2, 2]); - const offset = tf.tensor2d([0, 0, 0, 0], [2, 2]); - - const varianceEpsilon = .001; - - const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); - const gradX = tf.grad( - (x: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [1.414, 2.500, 0.816, 1.768]); - expect(gradX.shape).toEqual([2, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [-1.414, -2.500, -0.816, -1.768]); - expect(gradMean.shape).toEqual([2, 2]); - const gradVariance = tf.grad( - (variance: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose( - await gradVariance.data(), [-3.533, -4.686, -1.360, -2.762]); - expect(gradVariance.shape).toEqual([2, 2]); - const gradOffset = tf.grad( - (offset: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), [1, 1, 1, 1]); - expect(gradOffset.shape).toEqual([2, 2]); - const gradScale = tf.grad( - (scale: tf.Tensor2D) => tf.batchNorm2d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [7.069, 7.499, 8.164, 8.838]); - expect(gradScale.shape).toEqual([2, 2]); - }); - - it('gradient with clones', () => { - const x = tf.zeros([2, 2]); - const mean = tf.zeros([2, 2]); - const variance = tf.zeros([2, 2]); - const scale = tf.zeros([2, 2]); - const offset = tf.zeros([2, 2]); - - const varianceEpsilon = .001; - - const gradF = tf.grads( - (x: tf.Tensor2D, mean: tf.Tensor2D, variance: tf.Tensor2D, - offset: tf.Tensor2D, scale: tf.Tensor2D) => - tf.batchNorm2d( - x.clone(), mean.clone(), variance.clone(), offset.clone(), - scale.clone(), varianceEpsilon) - .clone()); - const [gradX, gradMean, gradVariance, gradOffset, gradScale] = - gradF([x, mean, variance, offset, scale]); - expect(gradX.shape).toEqual(x.shape); - expect(gradMean.shape).toEqual(mean.shape); - expect(gradVariance.shape).toEqual(variance.shape); - expect(gradOffset.shape).toEqual(offset.shape); - expect(gradScale.shape).toEqual(scale.shape); - }); - - it('batchnorm2D matches tensorflow, 3x3', async () => { - const x = tf.tensor2d( - [ - 0.3136892, 0.92389025, 0.594782, 0.05021042, 0.67545404, 0.93910035, - 0.13277993, 0.96474269, 0.88608916 - ], - [3, 3]); - const mean = tf.tensor1d([0.19526312, 0.74857256, 0.45166398]); - const variance = tf.tensor1d([0.22963001, 0.61521992, 0.46623685]); - const offset = tf.tensor1d([0.43098484, 0.77712237, 0.47916298]); - const scale = tf.tensor1d([0.62186907, 0.85673736, 0.19201061]); - const varianceEpsilon = .001; - - const result = - tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - 0.58433646, 0.96846228, 0.51936529, 0.24315402, 0.69732157, 0.61608542, - 0.35007446, 1.01304821, 0.60119441 - ]); - }); - - it('throws when passed x as a non-tensor', () => { - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - - expect(() => tf.batchNorm({} as tf.Tensor, mean, variance)) - .toThrowError(/Argument 'x' passed to 'batchNorm' must be a Tensor/); - }); - it('throws when passed mean as a non-tensor', () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const variance = tf.tensor1d([2, 3]); - - expect(() => tf.batchNorm(x, {} as tf.Tensor, variance)) - .toThrowError(/Argument 'mean' passed to 'batchNorm' must be a Tensor/); - }); - it('throws when passed variance as a non-tensor', () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const mean = tf.tensor1d([1, 2]); - - const e = /Argument 'variance' passed to 'batchNorm' must be a Tensor/; - expect(() => tf.batchNorm(x, mean, {} as tf.Tensor)).toThrowError(e); - }); - it('throws when passed scale as a non-tensor', () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const epsilon = .001; - - expect(() => tf.batchNorm(x, mean, variance, epsilon, {} as tf.Tensor)) - .toThrowError( - /Argument 'scale' passed to 'batchNorm' must be a Tensor/); - }); - it('throws when passed offset as a non-tensor', () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const epsilon = .001; - const scale = tf.tensor1d([0.62186907, 0.85673736, 0.19201061]); - - const e = /Argument 'offset' passed to 'batchNorm' must be a Tensor/; - expect( - () => tf.batchNorm(x, mean, variance, {} as tf.Tensor, scale, epsilon)) - .toThrowError(e); - }); - - it('accepts a tensor-like object', async () => { - const x = [[2, 4], [9, 23]]; - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const result = - tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('throws error when x is a string tensor', () => { - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const f = () => tf.batchNorm2d( - [['a', 'b'], ['c', 'd']], mean, variance, offset, scale, - varianceEpsilon); - expect(f).toThrowError( - /Argument 'x' passed to 'batchNorm' must be numeric/); - }); - - it('throws error when mean is a string tensor', () => { - const x = [[2, 4], [9, 23]]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const f = () => - tf.batchNorm2d(x, ['a', 'b'], variance, offset, scale, varianceEpsilon); - expect(f).toThrowError( - /Argument 'mean' passed to 'batchNorm' must be numeric/); - }); - - it('throws error when variance is a string tensor', () => { - const x = [[2, 4], [9, 23]]; - const mean = [1, 2]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const f = () => - tf.batchNorm2d(x, mean, ['a', 'b'], offset, scale, varianceEpsilon); - expect(f).toThrowError(/'variance' passed to 'batchNorm' must be numeric/); - }); - - it('throws error when scale is a string tensor', () => { - const x = [[2, 4], [9, 23]]; - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - - const varianceEpsilon = .001; - - const f = () => - tf.batchNorm2d(x, mean, variance, offset, ['a', 'b'], varianceEpsilon); - expect(f).toThrowError(/'scale' passed to 'batchNorm' must be numeric/); - }); - - it('throws error when offset is a string tensor', () => { - const x = [[2, 4], [9, 23]]; - const mean = [1, 2]; - const variance = [2, 3]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const f = () => - tf.batchNorm2d(x, mean, variance, ['a', 'b'], scale, varianceEpsilon); - expect(f).toThrowError(/'offset' passed to 'batchNorm' must be numeric/); - }); -}); - -describeWithFlags('deprecated batchNormalization', ALL_ENVS, () => { - beforeAll(() => { - // Silence deprecation warnings. - spyOn(console, 'warn'); - }); - - it('simple batchnorm2D, 2x2', async () => { - const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([3, 4]); - const scaleT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = tf.batchNormalization( - xT, meanT, varianceT, varianceEpsilon, scaleT, offsetT); - - const offset = await offsetT.array(); - const mean = await meanT.array(); - const variance = await varianceT.array(); - const scale = await scaleT.array(); - const x = await xT.array(); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - - const result2 = tf.batchNormalization2d( - xT, meanT, varianceT, varianceEpsilon, scaleT, offsetT); - expectArraysClose(await result.data(), await result2.data()); - }); -}); diff --git a/tfjs-core/src/ops/batchnorm3d.ts b/tfjs-core/src/ops/batchnorm3d.ts index 64638a4311e..e0edfbf9e93 100644 --- a/tfjs-core/src/ops/batchnorm3d.ts +++ b/tfjs-core/src/ops/batchnorm3d.ts @@ -36,7 +36,8 @@ import {warnDeprecation} from './batchnorm_util'; function batchNorm3d_( x: Tensor3D | TensorLike, mean: Tensor3D | Tensor1D | TensorLike, variance: Tensor3D | Tensor1D | TensorLike, - offset?: Tensor3D | Tensor1D | TensorLike, scale?: Tensor3D | Tensor1D | TensorLike, + offset?: Tensor3D | Tensor1D | TensorLike, + scale?: Tensor3D | Tensor1D | TensorLike, varianceEpsilon?: number): Tensor3D { const $x = convertToTensor(x, 'x', 'batchNorm'); const $mean = convertToTensor(mean, 'mean', 'batchNorm'); diff --git a/tfjs-core/src/ops/batchnorm3d_test.ts b/tfjs-core/src/ops/batchnorm3d_test.ts deleted file mode 100644 index 02473575604..00000000000 --- a/tfjs-core/src/ops/batchnorm3d_test.ts +++ /dev/null @@ -1,296 +0,0 @@ -/** - * @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 * as tf from '../index'; -import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysClose} from '../test_util'; - -describeWithFlags('batchNorm3D', ALL_ENVS, () => { - it('simple batchnorm3D, no offset or scale, 2x1x2', async () => { - const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const varianceEpsilon = .001; - - const result = tf.batchNorm3d( - xT, meanT, varianceT, undefined, undefined, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - expectArraysClose(await result.data(), [ - (x.get(0, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm3D, no offset, 2x1x2', async () => { - const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const scaleT = tf.tensor1d([4, 5]); - const varianceEpsilon = .001; - - const result = tf.batchNorm3d( - xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const scale = await scaleT.buffer(); - expectArraysClose(await result.data(), [ - (x.get(0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm3D, no scale, 2x1x2', async () => { - const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = tf.batchNorm3d( - xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const offset = await offsetT.buffer(); - expectArraysClose(await result.data(), [ - offset.get(0) + - (x.get(0, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(0, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon), - offset.get(0) + - (x.get(1, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(1, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm3D, 2x1x2', async () => { - const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([3, 4]); - const scaleT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = - tf.batchNorm3d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const offset = await offsetT.buffer(); - const scale = await scaleT.buffer(); - - expectArraysClose(await result.data(), [ - offset.get(0) + - (x.get(0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - offset.get(0) + - (x.get(1, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(1, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('accepts a tensor-like object', async () => { - const x = [[[2, 4]], [[9, 23]]]; // 2x1x2 - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const result = - tf.batchNorm3d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('batchnorm3D, x,mean,var,offset,scale are all 3D', async () => { - const shape: [number, number, number] = [2, 1, 2]; - const xT = tf.tensor3d([2, 4, 9, 23], shape); - const meanT = tf.tensor3d([1, 2, 3, 4], shape); - const varianceT = tf.tensor3d([2, 3, 4, 5], shape); - const offsetT = tf.tensor3d([3, 4, 5, 6], shape); - const scaleT = tf.tensor3d([4, 5, 6, 7], shape); - - const varianceEpsilon = .001; - - const result = - tf.batchNorm3d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); - - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const offset = await offsetT.buffer(); - const scale = await scaleT.buffer(); - expectArraysClose(await result.data(), [ - offset.get(0, 0, 0) + - (x.get(0, 0, 0) - mean.get(0, 0, 0)) * scale.get(0, 0, 0) / - Math.sqrt(variance.get(0, 0, 0) + varianceEpsilon), - offset.get(0, 0, 1) + - (x.get(0, 0, 1) - mean.get(0, 0, 1)) * scale.get(0, 0, 1) / - Math.sqrt(variance.get(0, 0, 1) + varianceEpsilon), - offset.get(1, 0, 0) + - (x.get(1, 0, 0) - mean.get(1, 0, 0)) * scale.get(1, 0, 0) / - Math.sqrt(variance.get(1, 0, 0) + varianceEpsilon), - offset.get(1, 0, 1) + - (x.get(1, 0, 1) - mean.get(1, 0, 1)) * scale.get(1, 0, 1) / - Math.sqrt(variance.get(1, 0, 1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm3D gradients, 2x1x2', async () => { - const x = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const offset = tf.tensor1d([3, 4]); - const scale = tf.tensor1d([2, 5]); - - const varianceEpsilon = .001; - - const dy = tf.tensor3d([1, 1, 1, 1], [2, 1, 2]); - const gradX = tf.grad( - (x: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); - expect(gradX.shape).toEqual([2, 1, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor1D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [-2.828, -5.773]); - expect(gradMean.shape).toEqual([2]); - const gradVariance = tf.grad( - (variance: tf.Tensor1D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); - expect(gradVariance.shape).toEqual([2]); - const gradOffset = tf.grad( - (offset: tf.Tensor1D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), [2, 2]); - expect(gradOffset.shape).toEqual([2]); - const gradScale = tf.grad( - (scale: tf.Tensor1D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [6.362, 13.277]); - expect(gradScale.shape).toEqual([2]); - }); - - it('batchnorm3D gradients, same shapes in x, mean and variance', async () => { - const x = tf.tensor3d([10, 20, 30, 40], [2, 1, 2]); - const mean = tf.tensor3d([0, 5, 10, 15], [2, 1, 2]); - const variance = tf.tensor3d([2, 4, 6, 8], [2, 1, 2]); - const scale = tf.tensor3d([2, 5, 2, 5], [2, 1, 2]); - const offset = tf.tensor3d([0, 0, 0, 0], [2, 1, 2]); - - const varianceEpsilon = .001; - - const dy = tf.tensor3d([1, 1, 1, 1], [2, 1, 2]); - const gradX = tf.grad( - (x: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [1.414, 2.500, 0.816, 1.768]); - expect(gradX.shape).toEqual([2, 1, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [-1.414, -2.500, -0.816, -1.768]); - expect(gradMean.shape).toEqual([2, 1, 2]); - const gradVariance = tf.grad( - (variance: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose( - await gradVariance.data(), [-3.533, -4.686, -1.360, -2.762]); - expect(gradVariance.shape).toEqual([2, 1, 2]); - const gradOffset = tf.grad( - (offset: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), [1, 1, 1, 1]); - expect(gradOffset.shape).toEqual([2, 1, 2]); - const gradScale = tf.grad( - (scale: tf.Tensor3D) => tf.batchNorm3d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [7.069, 7.499, 8.164, 8.838]); - expect(gradScale.shape).toEqual([2, 1, 2]); - }); - - it('batchnorm matches tensorflow, 2x3x3', async () => { - const x = tf.tensor3d( - [ - 0.49955603, 0.04158615, -1.09440524, 2.03854165, -0.61578344, - 2.87533573, 1.18105987, 0.807462, 1.87888837, 2.26563962, -0.37040935, - 1.35848753, -0.75347094, 0.15683117, 0.91925946, 0.34121279, - 0.92717143, 1.89683965 - ], - [2, 3, 3]); - const mean = tf.tensor1d([0.39745062, -0.48062894, 0.4847822]); - const variance = tf.tensor1d([0.32375343, 0.67117643, 1.08334653]); - const offset = tf.tensor1d([0.69398749, -1.29056387, 0.9429723]); - const scale = tf.tensor1d([-0.5607271, 0.9878457, 0.25181573]); - const varianceEpsilon = .001; - - const result = - tf.batchNorm3d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - 0.59352049, -0.66135202, 0.5610874, -0.92077015, -1.45341019, 1.52106473, - -0.07704776, 0.26144429, 1.28010017, -1.14422404, -1.15776136, 1.15425493, - 1.82644104, -0.52249442, 1.04803919, 0.74932291, 0.40568101, 1.2844412 - ]); - }); -}); diff --git a/tfjs-core/src/ops/batchnorm4d.ts b/tfjs-core/src/ops/batchnorm4d.ts index 5304cc7386d..e46bfa34125 100644 --- a/tfjs-core/src/ops/batchnorm4d.ts +++ b/tfjs-core/src/ops/batchnorm4d.ts @@ -36,7 +36,8 @@ import {warnDeprecation} from './batchnorm_util'; function batchNorm4d_( x: Tensor4D | TensorLike, mean: Tensor4D | Tensor1D | TensorLike, variance: Tensor4D | Tensor1D | TensorLike, - offset?: Tensor4D | Tensor1D | TensorLike, scale?: Tensor4D | Tensor1D | TensorLike, + offset?: Tensor4D | Tensor1D | TensorLike, + scale?: Tensor4D | Tensor1D | TensorLike, varianceEpsilon?: number): Tensor4D { const $x = convertToTensor(x, 'x', 'batchNorm'); const $mean = convertToTensor(mean, 'mean', 'batchNorm'); diff --git a/tfjs-core/src/ops/batchnorm4d_test.ts b/tfjs-core/src/ops/batchnorm4d_test.ts deleted file mode 100644 index 60464c5630e..00000000000 --- a/tfjs-core/src/ops/batchnorm4d_test.ts +++ /dev/null @@ -1,233 +0,0 @@ -/** - * @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 * as tf from '../index'; -import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; -import {expectArraysClose} from '../test_util'; - -describeWithFlags('batchNorm4D', ALL_ENVS, () => { - it('simple batchnorm4D, no offset or scale, 2x1x1x2', async () => { - const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const varianceEpsilon = .001; - - const result = tf.batchNorm4d( - xT, meanT, varianceT, undefined, undefined, varianceEpsilon); - - const x = await xT.array(); - const mean = await meanT.array(); - const variance = await varianceT.array(); - expectArraysClose(await result.data(), [ - (x[0][0][0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), - (x[0][0][0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon), - (x[1][0][0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), - (x[1][0][0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('simple batchnorm4D, no offset, 2x1x1x2', async () => { - const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const scaleT = tf.tensor1d([4, 5]); - const varianceEpsilon = .001; - - const result = tf.batchNorm4d( - xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const scale = await scaleT.buffer(); - - expectArraysClose(await result.data(), [ - (x.get(0, 0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(0, 0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - (x.get(1, 0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - (x.get(1, 0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm4D, no scale, 2x1x1x2', async () => { - const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = tf.batchNorm4d( - xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const offset = await offsetT.buffer(); - - expectArraysClose(await result.data(), [ - offset.get(0) + - (x.get(0, 0, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(0, 0, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon), - offset.get(0) + - (x.get(1, 0, 0, 0) - mean.get(0)) * 1 / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(1, 0, 0, 1) - mean.get(1)) * 1 / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('simple batchnorm4D, 2x1x1x2', async () => { - const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const meanT = tf.tensor1d([1, 2]); - const varianceT = tf.tensor1d([2, 3]); - const offsetT = tf.tensor1d([3, 4]); - const scaleT = tf.tensor1d([4, 5]); - - const varianceEpsilon = .001; - - const result = - tf.batchNorm4d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); - const x = await xT.buffer(); - const mean = await meanT.buffer(); - const variance = await varianceT.buffer(); - const scale = await scaleT.buffer(); - const offset = await offsetT.buffer(); - - expectArraysClose(await result.data(), [ - offset.get(0) + - (x.get(0, 0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(0, 0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon), - offset.get(0) + - (x.get(1, 0, 0, 0) - mean.get(0)) * scale.get(0) / - Math.sqrt(variance.get(0) + varianceEpsilon), - offset.get(1) + - (x.get(1, 0, 0, 1) - mean.get(1)) * scale.get(1) / - Math.sqrt(variance.get(1) + varianceEpsilon) - ]); - }); - - it('accepts a tensor-like object', async () => { - const x = [[[[2, 4]]], [[[9, 23]]]]; // 2x1x1x2 - const mean = [1, 2]; - const variance = [2, 3]; - const offset = [3, 4]; - const scale = [4, 5]; - - const varianceEpsilon = .001; - - const result = - tf.batchNorm4d(x, mean, variance, offset, scale, varianceEpsilon); - - expectArraysClose(await result.data(), [ - offset[0] + - (x[0][0][0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[0][0][0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon), - offset[0] + - (x[1][0][0][0] - mean[0]) * scale[0] / - Math.sqrt(variance[0] + varianceEpsilon), - offset[1] + - (x[1][0][0][1] - mean[1]) * scale[1] / - Math.sqrt(variance[1] + varianceEpsilon) - ]); - }); - - it('simple batchnorm4D gradients, 2x1x1x2', async () => { - const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); - const mean = tf.tensor1d([1, 2]); - const variance = tf.tensor1d([2, 3]); - const offset = tf.tensor1d([3, 4]); - const scale = tf.tensor1d([2, 5]); - - const varianceEpsilon = .001; - - const dy = tf.tensor4d([-1, -1, -1, -1], [2, 1, 1, 2]); - const gradX = tf.grad( - (x: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [-1.414, -2.887, -1.414, -2.887]); - expect(gradX.shape).toEqual([2, 1, 1, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor1D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [2.828, 5.773]); - expect(gradMean.shape).toEqual([2]); - const gradVariance = tf.grad( - (variance: tf.Tensor1D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose(await gradVariance.data(), [3.180, 11.060]); - expect(gradVariance.shape).toEqual([2]); - const gradOffset = tf.grad( - (offset: tf.Tensor1D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), await dy.sum([0, 1, 2]).data()); - expect(gradOffset.shape).toEqual([2]); - const gradScale = tf.grad( - (scale: tf.Tensor1D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [-6.362, -13.277]); - expect(gradScale.shape).toEqual([2]); - }); - - it('batchnorm4D gradients, same shapes in x, mean and variance', async () => { - const x = tf.tensor4d([10, 20, 30, 40], [2, 1, 1, 2]); - const mean = tf.tensor4d([0, 5, 10, 15], [2, 1, 1, 2]); - const variance = tf.tensor4d([2, 4, 6, 8], [2, 1, 1, 2]); - const scale = tf.tensor4d([2, 5, 2, 5], [2, 1, 1, 2]); - const offset = tf.tensor4d([0, 0, 0, 0], [2, 1, 1, 2]); - - const varianceEpsilon = .001; - - const dy = tf.tensor4d([-1, -1, -1, -1], [2, 1, 1, 2]); - const gradX = tf.grad( - (x: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(x, dy); - expectArraysClose(await gradX.data(), [-1.414, -2.500, -0.816, -1.768]); - expect(gradX.shape).toEqual([2, 1, 1, 2]); - const gradMean = tf.grad( - (mean: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); - expectArraysClose(await gradMean.data(), [1.414, 2.500, 0.816, 1.768]); - expect(gradMean.shape).toEqual([2, 1, 1, 2]); - const gradVariance = tf.grad( - (variance: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); - expectArraysClose(await gradVariance.data(), [3.533, 4.686, 1.360, 2.762]); - expect(gradVariance.shape).toEqual([2, 1, 1, 2]); - const gradOffset = tf.grad( - (offset: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); - expectArraysClose(await gradOffset.data(), await dy.data()); - expect(gradOffset.shape).toEqual([2, 1, 1, 2]); - const gradScale = tf.grad( - (scale: tf.Tensor4D) => tf.batchNorm4d( - x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); - expectArraysClose(await gradScale.data(), [-7.069, -7.499, -8.164, -8.838]); - expect(gradScale.shape).toEqual([2, 1, 1, 2]); - }); -}); diff --git a/tfjs-core/src/ops/batchnorm_test.ts b/tfjs-core/src/ops/batchnorm_test.ts new file mode 100644 index 00000000000..3b2e88b0f54 --- /dev/null +++ b/tfjs-core/src/ops/batchnorm_test.ts @@ -0,0 +1,950 @@ +/** + * @license + * Copyright 2017 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('batchNorm4D', ALL_ENVS, () => { + it('simple batchnorm4D, no offset or scale, 2x1x1x2', async () => { + const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const varianceEpsilon = .001; + + const result = tf.batchNorm4d( + xT, meanT, varianceT, undefined, undefined, varianceEpsilon); + + const x = await xT.array(); + const mean = await meanT.array(); + const variance = await varianceT.array(); + expectArraysClose(await result.data(), [ + (x[0][0][0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), + (x[0][0][0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon), + (x[1][0][0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), + (x[1][0][0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('simple batchnorm4D, no offset, 2x1x1x2', async () => { + const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const scaleT = tf.tensor1d([4, 5]); + const varianceEpsilon = .001; + + const result = tf.batchNorm4d( + xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const scale = await scaleT.buffer(); + + expectArraysClose(await result.data(), [ + (x.get(0, 0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm4D, no scale, 2x1x1x2', async () => { + const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = tf.batchNorm4d( + xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const offset = await offsetT.buffer(); + + expectArraysClose(await result.data(), [ + offset.get(0) + + (x.get(0, 0, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(0, 0, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon), + offset.get(0) + + (x.get(1, 0, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(1, 0, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm4D, 2x1x1x2', async () => { + const xT = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([3, 4]); + const scaleT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = + tf.batchNorm4d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const scale = await scaleT.buffer(); + const offset = await offsetT.buffer(); + + expectArraysClose(await result.data(), [ + offset.get(0) + + (x.get(0, 0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(0, 0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + offset.get(0) + + (x.get(1, 0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(1, 0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('accepts a tensor-like object', async () => { + const x = [[[[2, 4]]], [[[9, 23]]]]; // 2x1x1x2 + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const result = + tf.batchNorm4d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0][0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][0][0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0][0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][0][0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('simple batchnorm4D gradients, 2x1x1x2', async () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const offset = tf.tensor1d([3, 4]); + const scale = tf.tensor1d([2, 5]); + + const varianceEpsilon = .001; + + const dy = tf.tensor4d([-1, -1, -1, -1], [2, 1, 1, 2]); + const gradX = tf.grad( + (x: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [-1.414, -2.887, -1.414, -2.887]); + expect(gradX.shape).toEqual([2, 1, 1, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor1D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [2.828, 5.773]); + expect(gradMean.shape).toEqual([2]); + const gradVariance = tf.grad( + (variance: tf.Tensor1D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose(await gradVariance.data(), [3.180, 11.060]); + expect(gradVariance.shape).toEqual([2]); + const gradOffset = tf.grad( + (offset: tf.Tensor1D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), await dy.sum([0, 1, 2]).data()); + expect(gradOffset.shape).toEqual([2]); + const gradScale = tf.grad( + (scale: tf.Tensor1D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [-6.362, -13.277]); + expect(gradScale.shape).toEqual([2]); + }); + + it('batchnorm4D gradients, same shapes in x, mean and variance', async () => { + const x = tf.tensor4d([10, 20, 30, 40], [2, 1, 1, 2]); + const mean = tf.tensor4d([0, 5, 10, 15], [2, 1, 1, 2]); + const variance = tf.tensor4d([2, 4, 6, 8], [2, 1, 1, 2]); + const scale = tf.tensor4d([2, 5, 2, 5], [2, 1, 1, 2]); + const offset = tf.tensor4d([0, 0, 0, 0], [2, 1, 1, 2]); + + const varianceEpsilon = .001; + + const dy = tf.tensor4d([-1, -1, -1, -1], [2, 1, 1, 2]); + const gradX = tf.grad( + (x: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [-1.414, -2.500, -0.816, -1.768]); + expect(gradX.shape).toEqual([2, 1, 1, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [1.414, 2.500, 0.816, 1.768]); + expect(gradMean.shape).toEqual([2, 1, 1, 2]); + const gradVariance = tf.grad( + (variance: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose(await gradVariance.data(), [3.533, 4.686, 1.360, 2.762]); + expect(gradVariance.shape).toEqual([2, 1, 1, 2]); + const gradOffset = tf.grad( + (offset: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), await dy.data()); + expect(gradOffset.shape).toEqual([2, 1, 1, 2]); + const gradScale = tf.grad( + (scale: tf.Tensor4D) => tf.batchNorm4d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [-7.069, -7.499, -8.164, -8.838]); + expect(gradScale.shape).toEqual([2, 1, 1, 2]); + }); +}); + +describeWithFlags('batchNorm3D', ALL_ENVS, () => { + it('simple batchnorm3D, no offset or scale, 2x1x2', async () => { + const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const varianceEpsilon = .001; + + const result = tf.batchNorm3d( + xT, meanT, varianceT, undefined, undefined, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + expectArraysClose(await result.data(), [ + (x.get(0, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm3D, no offset, 2x1x2', async () => { + const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const scaleT = tf.tensor1d([4, 5]); + const varianceEpsilon = .001; + + const result = tf.batchNorm3d( + xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const scale = await scaleT.buffer(); + expectArraysClose(await result.data(), [ + (x.get(0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm3D, no scale, 2x1x2', async () => { + const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = tf.batchNorm3d( + xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const offset = await offsetT.buffer(); + expectArraysClose(await result.data(), [ + offset.get(0) + + (x.get(0, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(0, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon), + offset.get(0) + + (x.get(1, 0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(1, 0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm3D, 2x1x2', async () => { + const xT = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([3, 4]); + const scaleT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = + tf.batchNorm3d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const offset = await offsetT.buffer(); + const scale = await scaleT.buffer(); + + expectArraysClose(await result.data(), [ + offset.get(0) + + (x.get(0, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(0, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + offset.get(0) + + (x.get(1, 0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + offset.get(1) + + (x.get(1, 0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('accepts a tensor-like object', async () => { + const x = [[[2, 4]], [[9, 23]]]; // 2x1x2 + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const result = + tf.batchNorm3d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('batchnorm3D, x,mean,var,offset,scale are all 3D', async () => { + const shape: [number, number, number] = [2, 1, 2]; + const xT = tf.tensor3d([2, 4, 9, 23], shape); + const meanT = tf.tensor3d([1, 2, 3, 4], shape); + const varianceT = tf.tensor3d([2, 3, 4, 5], shape); + const offsetT = tf.tensor3d([3, 4, 5, 6], shape); + const scaleT = tf.tensor3d([4, 5, 6, 7], shape); + + const varianceEpsilon = .001; + + const result = + tf.batchNorm3d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const offset = await offsetT.buffer(); + const scale = await scaleT.buffer(); + expectArraysClose(await result.data(), [ + offset.get(0, 0, 0) + + (x.get(0, 0, 0) - mean.get(0, 0, 0)) * scale.get(0, 0, 0) / + Math.sqrt(variance.get(0, 0, 0) + varianceEpsilon), + offset.get(0, 0, 1) + + (x.get(0, 0, 1) - mean.get(0, 0, 1)) * scale.get(0, 0, 1) / + Math.sqrt(variance.get(0, 0, 1) + varianceEpsilon), + offset.get(1, 0, 0) + + (x.get(1, 0, 0) - mean.get(1, 0, 0)) * scale.get(1, 0, 0) / + Math.sqrt(variance.get(1, 0, 0) + varianceEpsilon), + offset.get(1, 0, 1) + + (x.get(1, 0, 1) - mean.get(1, 0, 1)) * scale.get(1, 0, 1) / + Math.sqrt(variance.get(1, 0, 1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm3D gradients, 2x1x2', async () => { + const x = tf.tensor3d([2, 4, 9, 23], [2, 1, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const offset = tf.tensor1d([3, 4]); + const scale = tf.tensor1d([2, 5]); + + const varianceEpsilon = .001; + + const dy = tf.tensor3d([1, 1, 1, 1], [2, 1, 2]); + const gradX = tf.grad( + (x: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); + expect(gradX.shape).toEqual([2, 1, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor1D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [-2.828, -5.773]); + expect(gradMean.shape).toEqual([2]); + const gradVariance = tf.grad( + (variance: tf.Tensor1D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); + expect(gradVariance.shape).toEqual([2]); + const gradOffset = tf.grad( + (offset: tf.Tensor1D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), [2, 2]); + expect(gradOffset.shape).toEqual([2]); + const gradScale = tf.grad( + (scale: tf.Tensor1D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [6.362, 13.277]); + expect(gradScale.shape).toEqual([2]); + }); + + it('batchnorm3D gradients, same shapes in x, mean and variance', async () => { + const x = tf.tensor3d([10, 20, 30, 40], [2, 1, 2]); + const mean = tf.tensor3d([0, 5, 10, 15], [2, 1, 2]); + const variance = tf.tensor3d([2, 4, 6, 8], [2, 1, 2]); + const scale = tf.tensor3d([2, 5, 2, 5], [2, 1, 2]); + const offset = tf.tensor3d([0, 0, 0, 0], [2, 1, 2]); + + const varianceEpsilon = .001; + + const dy = tf.tensor3d([1, 1, 1, 1], [2, 1, 2]); + const gradX = tf.grad( + (x: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [1.414, 2.500, 0.816, 1.768]); + expect(gradX.shape).toEqual([2, 1, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [-1.414, -2.500, -0.816, -1.768]); + expect(gradMean.shape).toEqual([2, 1, 2]); + const gradVariance = tf.grad( + (variance: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose( + await gradVariance.data(), [-3.533, -4.686, -1.360, -2.762]); + expect(gradVariance.shape).toEqual([2, 1, 2]); + const gradOffset = tf.grad( + (offset: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), [1, 1, 1, 1]); + expect(gradOffset.shape).toEqual([2, 1, 2]); + const gradScale = tf.grad( + (scale: tf.Tensor3D) => tf.batchNorm3d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [7.069, 7.499, 8.164, 8.838]); + expect(gradScale.shape).toEqual([2, 1, 2]); + }); + + it('batchnorm matches tensorflow, 2x3x3', async () => { + const x = tf.tensor3d( + [ + 0.49955603, 0.04158615, -1.09440524, 2.03854165, -0.61578344, + 2.87533573, 1.18105987, 0.807462, 1.87888837, 2.26563962, -0.37040935, + 1.35848753, -0.75347094, 0.15683117, 0.91925946, 0.34121279, + 0.92717143, 1.89683965 + ], + [2, 3, 3]); + const mean = tf.tensor1d([0.39745062, -0.48062894, 0.4847822]); + const variance = tf.tensor1d([0.32375343, 0.67117643, 1.08334653]); + const offset = tf.tensor1d([0.69398749, -1.29056387, 0.9429723]); + const scale = tf.tensor1d([-0.5607271, 0.9878457, 0.25181573]); + const varianceEpsilon = .001; + + const result = + tf.batchNorm3d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + 0.59352049, -0.66135202, 0.5610874, -0.92077015, -1.45341019, 1.52106473, + -0.07704776, 0.26144429, 1.28010017, -1.14422404, -1.15776136, 1.15425493, + 1.82644104, -0.52249442, 1.04803919, 0.74932291, 0.40568101, 1.2844412 + ]); + }); +}); + +describeWithFlags('batchNorm2D', ALL_ENVS, () => { + it('simple batchnorm2D, no offset or scale, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const varianceEpsilon = .001; + + const result = tf.batchNorm2d( + xT, meanT, varianceT, undefined, undefined, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + expectArraysClose(await result.data(), [ + (x.get(0, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0) - mean.get(0)) * 1 / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 1) - mean.get(1)) * 1 / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + it('simple batchnorm2D, no offset, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const scaleT = tf.tensor1d([4, 5]); + const varianceEpsilon = .001; + + const result = tf.batchNorm2d( + xT, meanT, varianceT, undefined, scaleT, varianceEpsilon); + + const x = await xT.buffer(); + const mean = await meanT.buffer(); + const variance = await varianceT.buffer(); + const scale = await scaleT.buffer(); + expectArraysClose(await result.data(), [ + (x.get(0, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(0, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon), + (x.get(1, 0) - mean.get(0)) * scale.get(0) / + Math.sqrt(variance.get(0) + varianceEpsilon), + (x.get(1, 1) - mean.get(1)) * scale.get(1) / + Math.sqrt(variance.get(1) + varianceEpsilon) + ]); + }); + + it('simple batchnorm2D, no scale, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = tf.batchNorm2d( + xT, meanT, varianceT, offsetT, undefined, varianceEpsilon); + + const offset = await offsetT.array(); + const mean = await meanT.array(); + const variance = await varianceT.array(); + const x = await xT.array(); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0] - mean[0]) * 1 / Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][1] - mean[1]) * 1 / Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('simple batchnorm2D, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([3, 4]); + const scaleT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = + tf.batchNorm2d(xT, meanT, varianceT, offsetT, scaleT, varianceEpsilon); + + const offset = await offsetT.array(); + const mean = await meanT.array(); + const variance = await varianceT.array(); + const scale = await scaleT.array(); + const x = await xT.array(); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('simple batchnorm2D gradients, 2x2', async () => { + const x = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const offset = tf.tensor1d([3, 4]); + const scale = tf.tensor1d([2, 5]); + + const varianceEpsilon = .001; + + const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); + const [gradX, gradMean, gradVariance, gradOffset, gradScale] = tf.grads( + (x: tf.Tensor2D, mean: tf.Tensor1D, variance: tf.Tensor1D, + offset: tf.Tensor1D, scale: tf.Tensor1D) => + tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon))( + [x, mean, variance, offset, scale], dy); + expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); + expect(gradX.shape).toEqual([2, 2]); + expectArraysClose(await gradMean.data(), [-2.828, -5.773]); + expect(gradMean.shape).toEqual([2]); + expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); + expect(gradVariance.shape).toEqual([2]); + expectArraysClose(await gradOffset.data(), [2, 2]); + expect(gradOffset.shape).toEqual([2]); + expectArraysClose(await gradScale.data(), [6.362, 13.277]); + expect(gradScale.shape).toEqual([2]); + }); + + it('gradient with clones batchnorm2D', async () => { + const x = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const offset = tf.tensor1d([3, 4]); + const scale = tf.tensor1d([2, 5]); + + const varianceEpsilon = .001; + + const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); + const [gradX, gradMean, gradVariance, gradOffset, gradScale] = tf.grads( + (x: tf.Tensor2D, mean: tf.Tensor1D, variance: tf.Tensor1D, + offset: tf.Tensor1D, scale: tf.Tensor1D) => + tf.batchNorm2d( + x.clone(), mean.clone(), variance.clone(), offset.clone(), + scale.clone(), varianceEpsilon) + .clone())([x, mean, variance, offset, scale], dy); + expectArraysClose(await gradX.data(), [1.414, 2.887, 1.414, 2.887]); + expect(gradX.shape).toEqual([2, 2]); + expectArraysClose(await gradMean.data(), [-2.828, -5.773]); + expect(gradMean.shape).toEqual([2]); + expectArraysClose(await gradVariance.data(), [-3.180, -11.060]); + expect(gradVariance.shape).toEqual([2]); + expectArraysClose(await gradOffset.data(), [2, 2]); + expect(gradOffset.shape).toEqual([2]); + expectArraysClose(await gradScale.data(), [6.362, 13.277]); + expect(gradScale.shape).toEqual([2]); + }); + + it('batchnorm2D gradients, same shapes in x, mean and variance', async () => { + const x = tf.tensor2d([10, 20, 30, 40], [2, 2]); + const mean = tf.tensor2d([0, 5, 10, 15], [2, 2]); + const variance = tf.tensor2d([2, 4, 6, 8], [2, 2]); + const scale = tf.tensor2d([2, 5, 2, 5], [2, 2]); + const offset = tf.tensor2d([0, 0, 0, 0], [2, 2]); + + const varianceEpsilon = .001; + + const dy = tf.tensor2d([1, 1, 1, 1], [2, 2]); + const gradX = tf.grad( + (x: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(x, dy); + expectArraysClose(await gradX.data(), [1.414, 2.500, 0.816, 1.768]); + expect(gradX.shape).toEqual([2, 2]); + const gradMean = tf.grad( + (mean: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(mean, dy); + expectArraysClose(await gradMean.data(), [-1.414, -2.500, -0.816, -1.768]); + expect(gradMean.shape).toEqual([2, 2]); + const gradVariance = tf.grad( + (variance: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(variance, dy); + expectArraysClose( + await gradVariance.data(), [-3.533, -4.686, -1.360, -2.762]); + expect(gradVariance.shape).toEqual([2, 2]); + const gradOffset = tf.grad( + (offset: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(offset, dy); + expectArraysClose(await gradOffset.data(), [1, 1, 1, 1]); + expect(gradOffset.shape).toEqual([2, 2]); + const gradScale = tf.grad( + (scale: tf.Tensor2D) => tf.batchNorm2d( + x, mean, variance, offset, scale, varianceEpsilon))(scale, dy); + expectArraysClose(await gradScale.data(), [7.069, 7.499, 8.164, 8.838]); + expect(gradScale.shape).toEqual([2, 2]); + }); + + it('gradient with clones', () => { + const x = tf.zeros([2, 2]); + const mean = tf.zeros([2, 2]); + const variance = tf.zeros([2, 2]); + const scale = tf.zeros([2, 2]); + const offset = tf.zeros([2, 2]); + + const varianceEpsilon = .001; + + const gradF = tf.grads( + (x: tf.Tensor2D, mean: tf.Tensor2D, variance: tf.Tensor2D, + offset: tf.Tensor2D, scale: tf.Tensor2D) => + tf.batchNorm2d( + x.clone(), mean.clone(), variance.clone(), offset.clone(), + scale.clone(), varianceEpsilon) + .clone()); + const [gradX, gradMean, gradVariance, gradOffset, gradScale] = + gradF([x, mean, variance, offset, scale]); + expect(gradX.shape).toEqual(x.shape); + expect(gradMean.shape).toEqual(mean.shape); + expect(gradVariance.shape).toEqual(variance.shape); + expect(gradOffset.shape).toEqual(offset.shape); + expect(gradScale.shape).toEqual(scale.shape); + }); + + it('batchnorm2D matches tensorflow, 3x3', async () => { + const x = tf.tensor2d( + [ + 0.3136892, 0.92389025, 0.594782, 0.05021042, 0.67545404, 0.93910035, + 0.13277993, 0.96474269, 0.88608916 + ], + [3, 3]); + const mean = tf.tensor1d([0.19526312, 0.74857256, 0.45166398]); + const variance = tf.tensor1d([0.22963001, 0.61521992, 0.46623685]); + const offset = tf.tensor1d([0.43098484, 0.77712237, 0.47916298]); + const scale = tf.tensor1d([0.62186907, 0.85673736, 0.19201061]); + const varianceEpsilon = .001; + + const result = + tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + 0.58433646, 0.96846228, 0.51936529, 0.24315402, 0.69732157, 0.61608542, + 0.35007446, 1.01304821, 0.60119441 + ]); + }); + + it('throws when passed x as a non-tensor', () => { + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + + expect(() => tf.batchNorm({} as tf.Tensor, mean, variance)) + .toThrowError(/Argument 'x' passed to 'batchNorm' must be a Tensor/); + }); + it('throws when passed mean as a non-tensor', () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const variance = tf.tensor1d([2, 3]); + + expect(() => tf.batchNorm(x, {} as tf.Tensor, variance)) + .toThrowError(/Argument 'mean' passed to 'batchNorm' must be a Tensor/); + }); + it('throws when passed variance as a non-tensor', () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const mean = tf.tensor1d([1, 2]); + + const e = /Argument 'variance' passed to 'batchNorm' must be a Tensor/; + expect(() => tf.batchNorm(x, mean, {} as tf.Tensor)).toThrowError(e); + }); + it('throws when passed scale as a non-tensor', () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const epsilon = .001; + + expect(() => tf.batchNorm(x, mean, variance, epsilon, {} as tf.Tensor)) + .toThrowError( + /Argument 'scale' passed to 'batchNorm' must be a Tensor/); + }); + it('throws when passed offset as a non-tensor', () => { + const x = tf.tensor4d([2, 4, 9, 23], [2, 1, 1, 2]); + const mean = tf.tensor1d([1, 2]); + const variance = tf.tensor1d([2, 3]); + const epsilon = .001; + const scale = tf.tensor1d([0.62186907, 0.85673736, 0.19201061]); + + const e = /Argument 'offset' passed to 'batchNorm' must be a Tensor/; + expect( + () => tf.batchNorm(x, mean, variance, {} as tf.Tensor, scale, epsilon)) + .toThrowError(e); + }); + + it('accepts a tensor-like object', async () => { + const x = [[2, 4], [9, 23]]; + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const result = + tf.batchNorm2d(x, mean, variance, offset, scale, varianceEpsilon); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + }); + + it('throws error when x is a string tensor', () => { + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const f = () => tf.batchNorm2d( + [['a', 'b'], ['c', 'd']], mean, variance, offset, scale, + varianceEpsilon); + expect(f).toThrowError( + /Argument 'x' passed to 'batchNorm' must be numeric/); + }); + + it('throws error when mean is a string tensor', () => { + const x = [[2, 4], [9, 23]]; + const variance = [2, 3]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const f = () => + tf.batchNorm2d(x, ['a', 'b'], variance, offset, scale, varianceEpsilon); + expect(f).toThrowError( + /Argument 'mean' passed to 'batchNorm' must be numeric/); + }); + + it('throws error when variance is a string tensor', () => { + const x = [[2, 4], [9, 23]]; + const mean = [1, 2]; + const offset = [3, 4]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const f = () => + tf.batchNorm2d(x, mean, ['a', 'b'], offset, scale, varianceEpsilon); + expect(f).toThrowError(/'variance' passed to 'batchNorm' must be numeric/); + }); + + it('throws error when scale is a string tensor', () => { + const x = [[2, 4], [9, 23]]; + const mean = [1, 2]; + const variance = [2, 3]; + const offset = [3, 4]; + + const varianceEpsilon = .001; + + const f = () => + tf.batchNorm2d(x, mean, variance, offset, ['a', 'b'], varianceEpsilon); + expect(f).toThrowError(/'scale' passed to 'batchNorm' must be numeric/); + }); + + it('throws error when offset is a string tensor', () => { + const x = [[2, 4], [9, 23]]; + const mean = [1, 2]; + const variance = [2, 3]; + const scale = [4, 5]; + + const varianceEpsilon = .001; + + const f = () => + tf.batchNorm2d(x, mean, variance, ['a', 'b'], scale, varianceEpsilon); + expect(f).toThrowError(/'offset' passed to 'batchNorm' must be numeric/); + }); +}); + +describeWithFlags('deprecated batchNormalization', ALL_ENVS, () => { + beforeAll(() => { + // Silence deprecation warnings. + spyOn(console, 'warn'); + }); + + it('simple batchnorm2D, 2x2', async () => { + const xT = tf.tensor2d([2, 4, 9, 23], [2, 2]); + const meanT = tf.tensor1d([1, 2]); + const varianceT = tf.tensor1d([2, 3]); + const offsetT = tf.tensor1d([3, 4]); + const scaleT = tf.tensor1d([4, 5]); + + const varianceEpsilon = .001; + + const result = tf.batchNormalization( + xT, meanT, varianceT, varianceEpsilon, scaleT, offsetT); + + const offset = await offsetT.array(); + const mean = await meanT.array(); + const variance = await varianceT.array(); + const scale = await scaleT.array(); + const x = await xT.array(); + + expectArraysClose(await result.data(), [ + offset[0] + + (x[0][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[0][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon), + offset[0] + + (x[1][0] - mean[0]) * scale[0] / + Math.sqrt(variance[0] + varianceEpsilon), + offset[1] + + (x[1][1] - mean[1]) * scale[1] / + Math.sqrt(variance[1] + varianceEpsilon) + ]); + + const result2 = tf.batchNormalization2d( + xT, meanT, varianceT, varianceEpsilon, scaleT, offsetT); + expectArraysClose(await result.data(), await result2.data()); + }); +}); diff --git a/tfjs-core/src/public/chained_ops/batchnorm.ts b/tfjs-core/src/public/chained_ops/batchnorm.ts index 0e26af18d97..ce621c8eb04 100644 --- a/tfjs-core/src/public/chained_ops/batchnorm.ts +++ b/tfjs-core/src/public/chained_ops/batchnorm.ts @@ -37,4 +37,4 @@ Tensor.prototype.batchNorm = function ( scale?: Tensor | Tensor1D | TensorLike, varianceEpsilon?: number): Tensor { return batchNorm(this, mean, variance, offset, scale, varianceEpsilon); -} +}; diff --git a/tfjs-core/src/tests.ts b/tfjs-core/src/tests.ts index 7ddf226263f..ae9143836b5 100644 --- a/tfjs-core/src/tests.ts +++ b/tfjs-core/src/tests.ts @@ -40,9 +40,7 @@ import './kernel_registry_test'; import './ops/arithmetic_test'; import './ops/array_ops_test'; import './ops/axis_util_test'; -import './ops/batchnorm2d_test'; -import './ops/batchnorm3d_test'; -import './ops/batchnorm4d_test'; +import './ops/batchnorm_test'; import './ops/binary_ops_test'; import './ops/boolean_mask_test'; import './ops/broadcast_to_test'; From b94a5a20b9d9a3de50537b1ae6285e5dc8694508 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 30 Mar 2020 14:13:22 -0700 Subject: [PATCH 3/8] . --- .vscode/settings.json | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index 997d33d2d42..2a84b4c07db 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -18,8 +18,7 @@ "editor.tabSize": 2, "editor.insertSpaces": true, "[typescript]": { - "editor.formatOnSave": true, - "editor.defaultFormatter": "xaver.clang-format" + "editor.formatOnSave": true }, "[javascript]": { "editor.formatOnSave": true From 11803937a28a2a3b1764cdfbf70a606178f5b7cb Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 30 Mar 2020 16:22:06 -0700 Subject: [PATCH 4/8] Refactor. --- .../src/gradients/BatchNormalization_grad.ts | 104 +++++++++++++++ tfjs-core/src/kernel_names.ts | 7 ++ tfjs-core/src/ops/batchnorm.ts | 118 ++++-------------- tfjs-core/src/ops/batchnorm_util.ts | 25 +++- tfjs-core/src/register_all_gradients.ts | 7 +- 5 files changed, 157 insertions(+), 104 deletions(-) create mode 100644 tfjs-core/src/gradients/BatchNormalization_grad.ts diff --git a/tfjs-core/src/gradients/BatchNormalization_grad.ts b/tfjs-core/src/gradients/BatchNormalization_grad.ts new file mode 100644 index 00000000000..d98c49da660 --- /dev/null +++ b/tfjs-core/src/gradients/BatchNormalization_grad.ts @@ -0,0 +1,104 @@ +/** + * @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 {BatchNormalization, BatchNormalizationAttrs} from '../kernel_names'; +import {GradConfig, NamedAttrMap} from '../kernel_registry'; +import {xAs4D} from '../ops/batchnorm_util'; +import {getReductionAxes} from '../ops/broadcast_util'; +import {scalar} from '../ops/tensor_ops'; +import {tile} from '../ops/tile'; +import {rsqrt} from '../ops/unary_ops'; +import {Tensor, Tensor4D} from '../tensor'; +import {Rank, ShapeMap} from '../types'; + +export const batchNormalizationGradConfig: GradConfig = { + kernelName: BatchNormalization, + inputsToSave: ['x', 'mean', 'variance', 'scale'], + gradFunc: ( + dy: Tensor, saved: Tensor[], attrs: NamedAttrMap) => { + const batchNormalizationAttrs: BatchNormalizationAttrs = + attrs as {} as BatchNormalizationAttrs; + const {varianceEpsilon} = batchNormalizationAttrs; + const [$x, $mean, $variance, $scale] = saved; + + const x4D: Tensor4D = xAs4D($x); + + const scaleValue = $scale == null ? scalar(1) : $scale; + const reductionAxes = getReductionAxes($mean.shape, x4D.shape); + const tileShape: number[] = []; + if ($mean.rank === 1) { + for (let i = 0; i < x4D.shape.length - 1; ++i) { + tileShape.push(x4D.shape[i]); + } + tileShape.push(1); + } + + const xMinusMean = $x.sub($mean); + const dyTimesScaleValue = dy.mul(scaleValue); + const oneOverSqrtVariance = rsqrt($variance.add(scalar(varianceEpsilon))); + const minusHalfRCube = oneOverSqrtVariance.mul(oneOverSqrtVariance) + .mul(oneOverSqrtVariance) + .mul(scalar(-0.5)); + + const derX = () => { + if ($mean.rank === 1) { + return dy + .mul(tile( + oneOverSqrtVariance.as4D(1, 1, 1, $mean.shape[0]), tileShape)) + .mul(scaleValue) + .reshape($x.shape); + } else { + return dy.mul(oneOverSqrtVariance).mul(scaleValue).reshape($x.shape); + } + }; + const derMean = () => { + let meanDer = oneOverSqrtVariance.mul(scalar(-1)).mul(dyTimesScaleValue); + if ($mean.rank === 1) { + meanDer = meanDer.sum(reductionAxes); + } + return meanDer.reshape($mean.shape as ShapeMap[R]); + }; + const derVariance = () => { + let varianceDer = minusHalfRCube.mul(xMinusMean).mul(dyTimesScaleValue); + if ($mean.rank === 1) { + varianceDer = varianceDer.sum(reductionAxes); + } + return varianceDer.reshape($mean.shape as ShapeMap[R]); + }; + const derScale = () => { + const xMinusMean2TimesRsqrt = xMinusMean.mul(oneOverSqrtVariance); + let scaleDer = dy.mul(xMinusMean2TimesRsqrt); + if ($mean.rank === 1) { + scaleDer = scaleDer.sum(reductionAxes); + } + return scaleDer.reshape($mean.shape as ShapeMap[R]); + }; + const derOffset = () => { + let offsetDer = dy; + if ($mean.rank === 1) { + offsetDer = offsetDer.sum(reductionAxes); + } + return offsetDer.reshape($mean.shape as ShapeMap[R]); + }; + return { + x: derX, + mean: derMean, + variance: derVariance, + scale: derScale, + offset: derOffset + }; + } +}; diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts index 239d9794a46..f66124cc33e 100644 --- a/tfjs-core/src/kernel_names.ts +++ b/tfjs-core/src/kernel_names.ts @@ -21,6 +21,13 @@ import {NamedTensorInfoMap} from './kernel_registry'; import {PixelData} from './types'; +export const BatchNormalization = 'BatchNormalization'; +export type BatchNormalizationInputs = + Pick; +export interface BatchNormalizationAttrs { + varianceEpsilon: number; +} + export type BinaryInputs = Pick; export const Div = 'Div'; diff --git a/tfjs-core/src/ops/batchnorm.ts b/tfjs-core/src/ops/batchnorm.ts index ecbb7b42323..45cd505603f 100644 --- a/tfjs-core/src/ops/batchnorm.ts +++ b/tfjs-core/src/ops/batchnorm.ts @@ -15,18 +15,17 @@ * ============================================================================= */ -import {ENGINE} from '../engine'; +import {ENGINE, ForwardFunc} from '../engine'; +import {BatchNormalizationAttrs, BatchNormalizationInputs} from '../kernel_names'; +import {NamedAttrMap} from '../kernel_registry'; import {Tensor, Tensor1D, Tensor4D} from '../tensor'; +import {NamedTensorMap} from '../tensor_types'; import {convertToTensor} from '../tensor_util_env'; -import {Rank, ShapeMap, TensorLike} from '../types'; +import {Rank, TensorLike} from '../types'; import * as util from '../util'; -import {warnDeprecation} from './batchnorm_util'; -import {getReductionAxes} from './broadcast_util'; +import {warnDeprecation, xAs4D} from './batchnorm_util'; import {op} from './operation'; -import {scalar} from './tensor_ops'; -import {tile} from './tile'; -import {rsqrt} from './unary_ops'; /** * @deprecated Please use `tf.batchNorm` instead and note the positional @@ -100,101 +99,26 @@ function batchNorm_( () => 'Batch normalization gradient requires mean and scale to have ' + 'equal ranks.'); - let x4D: Tensor4D; - if ($x.rank === 0 || $x.rank === 1) { - x4D = $x.as4D(1, 1, 1, $x.size); - } else if ($x.rank === 2) { - x4D = $x.as4D(1, 1, $x.shape[0], $x.shape[1]); - } else if ($x.rank === 3) { - x4D = $x.as4D(1, $x.shape[0], $x.shape[1], $x.shape[2]); - } else { - x4D = $x as Tensor4D; - } - - const der = (dy: Tensor, saved: Tensor[]) => { - type Saved = [ - Tensor, Tensor| Tensor1D, Tensor| Tensor1D, Tensor| Tensor1D - ]; - const [$x, $mean, $variance, $scale] = saved as Saved; - const scaleValue = $scale == null ? scalar(1) : $scale; - const reductionAxes = getReductionAxes($mean.shape, x4D.shape); - const tileShape: number[] = []; - if ($mean.rank === 1) { - for (let i = 0; i < x4D.shape.length - 1; ++i) { - tileShape.push(x4D.shape[i]); - } - tileShape.push(1); - } - - const xMinusMean = $x.sub($mean); - const dyTimesScaleValue = dy.mul(scaleValue); - const oneOverSqrtVariance = rsqrt($variance.add(scalar(varianceEpsilon))); - const minusHalfRCube = oneOverSqrtVariance.mul(oneOverSqrtVariance) - .mul(oneOverSqrtVariance) - .mul(scalar(-0.5)); + const x4D: Tensor4D = xAs4D($x); - const derX = () => { - if ($mean.rank === 1) { - return dy - .mul(tile( - oneOverSqrtVariance.as4D(1, 1, 1, $mean.shape[0]), tileShape)) - .mul(scaleValue) - .reshape($x.shape); - } else { - return dy.mul(oneOverSqrtVariance).mul(scaleValue).reshape($x.shape); - } - }; - const derMean = () => { - let meanDer = oneOverSqrtVariance.mul(scalar(-1)).mul(dyTimesScaleValue); - if ($mean.rank === 1) { - meanDer = meanDer.sum(reductionAxes); - } - return meanDer.reshape($mean.shape as ShapeMap[R]); - }; - const derVariance = () => { - let varianceDer = minusHalfRCube.mul(xMinusMean).mul(dyTimesScaleValue); - if ($mean.rank === 1) { - varianceDer = varianceDer.sum(reductionAxes); - } - return varianceDer.reshape($mean.shape as ShapeMap[R]); - }; - const derScale = () => { - const xMinusMean2TimesRsqrt = xMinusMean.mul(oneOverSqrtVariance); - let scaleDer = dy.mul(xMinusMean2TimesRsqrt); - if ($mean.rank === 1) { - scaleDer = scaleDer.sum(reductionAxes); - } - return scaleDer.reshape($mean.shape as ShapeMap[R]); - }; - const derOffset = () => { - let offsetDer = dy; - if ($mean.rank === 1) { - offsetDer = offsetDer.sum(reductionAxes); - } - return offsetDer.reshape($mean.shape as ShapeMap[R]); - }; - return { - x: derX, - mean: derMean, - variance: derVariance, - scale: derScale, - offset: derOffset - }; + const forward: ForwardFunc = (backend, save) => { + const res = backend.batchNormalization( + x4D, batchnormReshape4D($mean), batchnormReshape4D($variance), + varianceEpsilon, batchnormReshape4D($scale), + batchnormReshape4D($offset)); + save([$x, $mean, $variance, $scale]); + return res; }; - const inputsToSave = [$x, $mean, $variance, $scale]; + const inputs: BatchNormalizationInputs = + {x: $x, scale: $scale, offset: $offset, mean: $mean, variance: $variance}; + + const attrs: BatchNormalizationAttrs = {varianceEpsilon}; const res = ENGINE.runKernelFunc( - (backend, save) => { - const res = backend.batchNormalization( - x4D, batchnormReshape4D($mean), batchnormReshape4D($variance), - varianceEpsilon, batchnormReshape4D($scale), - batchnormReshape4D($offset)); - save([$x, $mean, $variance, $scale]); - return res; - }, - {x: $x, mean: $mean, variance: $variance, scale: $scale, offset: $offset}, - der, 'BatchNormalization', {varianceEpsilon}, inputsToSave); + forward, inputs as {} as NamedTensorMap, null /* gradient */, + 'BatchNormalization', attrs as {} as NamedAttrMap); + return res.reshape($x.shape); } diff --git a/tfjs-core/src/ops/batchnorm_util.ts b/tfjs-core/src/ops/batchnorm_util.ts index b2b4e2f74c3..715a28579d3 100644 --- a/tfjs-core/src/ops/batchnorm_util.ts +++ b/tfjs-core/src/ops/batchnorm_util.ts @@ -15,10 +15,27 @@ * ============================================================================= */ import {deprecationWarn} from '../globals'; +import {Tensor, Tensor4D} from '../tensor'; +import {Rank} from '../types'; -export function warnDeprecation() { +export function warnDeprecation(): void { deprecationWarn( - 'tf.batchNormalization() is going away. ' + - 'Use tf.batchNorm() instead, and note the positional argument change ' + - 'of scale, offset, and varianceEpsilon'); + 'tf.batchNormalization() is going away. ' + + 'Use tf.batchNorm() instead, and note the positional argument change ' + + 'of scale, offset, and varianceEpsilon'); +} + +export function xAs4D(x: Tensor) { + let x4D: Tensor4D; + if (x.rank === 0 || x.rank === 1) { + x4D = x.as4D(1, 1, 1, x.size); + } else if (x.rank === 2) { + x4D = x.as4D(1, 1, x.shape[0], x.shape[1]); + } else if (x.rank === 3) { + x4D = x.as4D(1, x.shape[0], x.shape[1], x.shape[2]); + } else { + x4D = x as Tensor4D; + } + + return x4D; } diff --git a/tfjs-core/src/register_all_gradients.ts b/tfjs-core/src/register_all_gradients.ts index a35522b2998..8db20b82768 100644 --- a/tfjs-core/src/register_all_gradients.ts +++ b/tfjs-core/src/register_all_gradients.ts @@ -14,6 +14,7 @@ * limitations under the License. * ============================================================================= */ +import {batchNormalizationGradConfig} from './gradients/BatchNormalization_grad'; import {broadcastToGradConfig} from './gradients/BroadcastTo_grad'; import {divGradConfig} from './gradients/Div_grad'; import {identityGradConfig} from './gradients/Identity_grad'; @@ -28,9 +29,9 @@ import {registerGradient} from './kernel_registry'; // Export all kernel configs here so that the package can auto register them const gradConfigs: GradConfig[] = [ - divGradConfig, squareGradConfig, squaredDifferenceGradConfig, - broadcastToGradConfig, identityGradConfig, tileGradConfig, oneHotGradConfig, - transposeGradConfig, padV2GradConfig + batchNormalizationGradConfig, divGradConfig, squareGradConfig, + squaredDifferenceGradConfig, broadcastToGradConfig, identityGradConfig, + tileGradConfig, oneHotGradConfig, transposeGradConfig, padV2GradConfig ]; for (const gradientConfig of gradConfigs) { From 218d36eda201090860dc12a6336bebaea85552d2 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 30 Mar 2020 16:26:09 -0700 Subject: [PATCH 5/8] Change to without $. --- .../src/gradients/BatchNormalization_grad.ts | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/tfjs-core/src/gradients/BatchNormalization_grad.ts b/tfjs-core/src/gradients/BatchNormalization_grad.ts index d98c49da660..59e6e4b2877 100644 --- a/tfjs-core/src/gradients/BatchNormalization_grad.ts +++ b/tfjs-core/src/gradients/BatchNormalization_grad.ts @@ -32,66 +32,66 @@ export const batchNormalizationGradConfig: GradConfig = { const batchNormalizationAttrs: BatchNormalizationAttrs = attrs as {} as BatchNormalizationAttrs; const {varianceEpsilon} = batchNormalizationAttrs; - const [$x, $mean, $variance, $scale] = saved; + const [x, mean, variance, scale] = saved; - const x4D: Tensor4D = xAs4D($x); + const x4D: Tensor4D = xAs4D(x); - const scaleValue = $scale == null ? scalar(1) : $scale; - const reductionAxes = getReductionAxes($mean.shape, x4D.shape); + const scaleValue = scale == null ? scalar(1) : scale; + const reductionAxes = getReductionAxes(mean.shape, x4D.shape); const tileShape: number[] = []; - if ($mean.rank === 1) { + if (mean.rank === 1) { for (let i = 0; i < x4D.shape.length - 1; ++i) { tileShape.push(x4D.shape[i]); } tileShape.push(1); } - const xMinusMean = $x.sub($mean); + const xMinusMean = x.sub(mean); const dyTimesScaleValue = dy.mul(scaleValue); - const oneOverSqrtVariance = rsqrt($variance.add(scalar(varianceEpsilon))); + const oneOverSqrtVariance = rsqrt(variance.add(scalar(varianceEpsilon))); const minusHalfRCube = oneOverSqrtVariance.mul(oneOverSqrtVariance) .mul(oneOverSqrtVariance) .mul(scalar(-0.5)); const derX = () => { - if ($mean.rank === 1) { + if (mean.rank === 1) { return dy .mul(tile( - oneOverSqrtVariance.as4D(1, 1, 1, $mean.shape[0]), tileShape)) + oneOverSqrtVariance.as4D(1, 1, 1, mean.shape[0]), tileShape)) .mul(scaleValue) - .reshape($x.shape); + .reshape(x.shape); } else { - return dy.mul(oneOverSqrtVariance).mul(scaleValue).reshape($x.shape); + return dy.mul(oneOverSqrtVariance).mul(scaleValue).reshape(x.shape); } }; const derMean = () => { let meanDer = oneOverSqrtVariance.mul(scalar(-1)).mul(dyTimesScaleValue); - if ($mean.rank === 1) { + if (mean.rank === 1) { meanDer = meanDer.sum(reductionAxes); } - return meanDer.reshape($mean.shape as ShapeMap[R]); + return meanDer.reshape(mean.shape as ShapeMap[R]); }; const derVariance = () => { let varianceDer = minusHalfRCube.mul(xMinusMean).mul(dyTimesScaleValue); - if ($mean.rank === 1) { + if (mean.rank === 1) { varianceDer = varianceDer.sum(reductionAxes); } - return varianceDer.reshape($mean.shape as ShapeMap[R]); + return varianceDer.reshape(mean.shape as ShapeMap[R]); }; const derScale = () => { const xMinusMean2TimesRsqrt = xMinusMean.mul(oneOverSqrtVariance); let scaleDer = dy.mul(xMinusMean2TimesRsqrt); - if ($mean.rank === 1) { + if (mean.rank === 1) { scaleDer = scaleDer.sum(reductionAxes); } - return scaleDer.reshape($mean.shape as ShapeMap[R]); + return scaleDer.reshape(mean.shape as ShapeMap[R]); }; const derOffset = () => { let offsetDer = dy; - if ($mean.rank === 1) { + if (mean.rank === 1) { offsetDer = offsetDer.sum(reductionAxes); } - return offsetDer.reshape($mean.shape as ShapeMap[R]); + return offsetDer.reshape(mean.shape as ShapeMap[R]); }; return { x: derX, From 615ffb6715c303341014a5f011754e1ef324cb9d Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 30 Mar 2020 16:33:22 -0700 Subject: [PATCH 6/8] Refactor function name. --- tfjs-core/src/ops/batchnorm.ts | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tfjs-core/src/ops/batchnorm.ts b/tfjs-core/src/ops/batchnorm.ts index 45cd505603f..d2475908d5e 100644 --- a/tfjs-core/src/ops/batchnorm.ts +++ b/tfjs-core/src/ops/batchnorm.ts @@ -103,9 +103,8 @@ function batchNorm_( const forward: ForwardFunc = (backend, save) => { const res = backend.batchNormalization( - x4D, batchnormReshape4D($mean), batchnormReshape4D($variance), - varianceEpsilon, batchnormReshape4D($scale), - batchnormReshape4D($offset)); + x4D, as1DOr4D($mean), as1DOr4D($variance), varianceEpsilon, + as1DOr4D($scale), as1DOr4D($offset)); save([$x, $mean, $variance, $scale]); return res; }; @@ -122,7 +121,7 @@ function batchNorm_( return res.reshape($x.shape); } -function batchnormReshape4D(x: Tensor): Tensor4D|Tensor1D { +function as1DOr4D(x: Tensor): Tensor4D|Tensor1D { if (x == null) { return null; } From aa3311ccf2d5f7909b21b7373d4d4adf045df539 Mon Sep 17 00:00:00 2001 From: Na Li Date: Tue, 31 Mar 2020 11:36:40 -0700 Subject: [PATCH 7/8] Change to use FusedBatchNorm as kernel name. Also change all op to be called directly in grad. --- ...ization_grad.ts => FusedBatchNorm_grad.ts} | 65 ++++++++++--------- tfjs-core/src/kernel_names.ts | 6 +- tfjs-core/src/ops/batchnorm.ts | 15 +++-- tfjs-core/src/ops/batchnorm2d.ts | 54 +++++++-------- tfjs-core/src/ops/batchnorm3d.ts | 55 ++++++++-------- tfjs-core/src/ops/batchnorm4d.ts | 55 ++++++++-------- tfjs-core/src/register_all_gradients.ts | 8 +-- 7 files changed, 135 insertions(+), 123 deletions(-) rename tfjs-core/src/gradients/{BatchNormalization_grad.ts => FusedBatchNorm_grad.ts} (56%) diff --git a/tfjs-core/src/gradients/BatchNormalization_grad.ts b/tfjs-core/src/gradients/FusedBatchNorm_grad.ts similarity index 56% rename from tfjs-core/src/gradients/BatchNormalization_grad.ts rename to tfjs-core/src/gradients/FusedBatchNorm_grad.ts index 59e6e4b2877..630fe68e2fc 100644 --- a/tfjs-core/src/gradients/BatchNormalization_grad.ts +++ b/tfjs-core/src/gradients/FusedBatchNorm_grad.ts @@ -14,23 +14,25 @@ * limitations under the License. * ============================================================================= */ -import {BatchNormalization, BatchNormalizationAttrs} from '../kernel_names'; +import {FusedBatchNorm, FusedBatchNormAttrs} from '../kernel_names'; import {GradConfig, NamedAttrMap} from '../kernel_registry'; import {xAs4D} from '../ops/batchnorm_util'; import {getReductionAxes} from '../ops/broadcast_util'; +import {add, mul, reshape, sub} from '../ops/ops'; +import {sum} from '../ops/reduction_ops'; import {scalar} from '../ops/tensor_ops'; import {tile} from '../ops/tile'; import {rsqrt} from '../ops/unary_ops'; import {Tensor, Tensor4D} from '../tensor'; import {Rank, ShapeMap} from '../types'; -export const batchNormalizationGradConfig: GradConfig = { - kernelName: BatchNormalization, +export const fusedBatchNormGradConfig: GradConfig = { + kernelName: FusedBatchNorm, inputsToSave: ['x', 'mean', 'variance', 'scale'], gradFunc: ( dy: Tensor, saved: Tensor[], attrs: NamedAttrMap) => { - const batchNormalizationAttrs: BatchNormalizationAttrs = - attrs as {} as BatchNormalizationAttrs; + const batchNormalizationAttrs: FusedBatchNormAttrs = + attrs as {} as FusedBatchNormAttrs; const {varianceEpsilon} = batchNormalizationAttrs; const [x, mean, variance, scale] = saved; @@ -46,52 +48,57 @@ export const batchNormalizationGradConfig: GradConfig = { tileShape.push(1); } - const xMinusMean = x.sub(mean); - const dyTimesScaleValue = dy.mul(scaleValue); - const oneOverSqrtVariance = rsqrt(variance.add(scalar(varianceEpsilon))); - const minusHalfRCube = oneOverSqrtVariance.mul(oneOverSqrtVariance) - .mul(oneOverSqrtVariance) - .mul(scalar(-0.5)); + const xMinusMean = sub(x, mean); + const dyTimesScaleValue = mul(dy, scaleValue); + const oneOverSqrtVariance = rsqrt(add(variance, scalar(varianceEpsilon))); + const minusHalfRCube = mul( + mul(mul(oneOverSqrtVariance, oneOverSqrtVariance), oneOverSqrtVariance), + scalar(-0.5)); const derX = () => { if (mean.rank === 1) { - return dy - .mul(tile( - oneOverSqrtVariance.as4D(1, 1, 1, mean.shape[0]), tileShape)) - .mul(scaleValue) - .reshape(x.shape); + return reshape( + mul(mul(dy, + tile( + oneOverSqrtVariance.as4D(1, 1, 1, mean.shape[0]), + tileShape)), + scaleValue), + x.shape); } else { - return dy.mul(oneOverSqrtVariance).mul(scaleValue).reshape(x.shape); + return reshape(mul(mul(dy, oneOverSqrtVariance), scaleValue), x.shape); } }; const derMean = () => { - let meanDer = oneOverSqrtVariance.mul(scalar(-1)).mul(dyTimesScaleValue); + let meanDer = + mul(mul(oneOverSqrtVariance, scalar(-1)), dyTimesScaleValue); if (mean.rank === 1) { - meanDer = meanDer.sum(reductionAxes); + meanDer = sum(meanDer, reductionAxes); } - return meanDer.reshape(mean.shape as ShapeMap[R]); + return reshape(meanDer, mean.shape as ShapeMap[R]); }; const derVariance = () => { - let varianceDer = minusHalfRCube.mul(xMinusMean).mul(dyTimesScaleValue); + let varianceDer = mul(mul(minusHalfRCube, xMinusMean), dyTimesScaleValue); + if (mean.rank === 1) { - varianceDer = varianceDer.sum(reductionAxes); + varianceDer = sum(varianceDer, reductionAxes); } - return varianceDer.reshape(mean.shape as ShapeMap[R]); + return reshape(varianceDer, mean.shape as ShapeMap[R]); }; const derScale = () => { - const xMinusMean2TimesRsqrt = xMinusMean.mul(oneOverSqrtVariance); - let scaleDer = dy.mul(xMinusMean2TimesRsqrt); + const xMinusMean2TimesRsqrt = mul(xMinusMean, oneOverSqrtVariance); + + let scaleDer = mul(dy, xMinusMean2TimesRsqrt); if (mean.rank === 1) { - scaleDer = scaleDer.sum(reductionAxes); + scaleDer = sum(scaleDer, reductionAxes); } - return scaleDer.reshape(mean.shape as ShapeMap[R]); + return reshape(scaleDer, mean.shape as ShapeMap[R]); }; const derOffset = () => { let offsetDer = dy; if (mean.rank === 1) { - offsetDer = offsetDer.sum(reductionAxes); + offsetDer = sum(offsetDer, reductionAxes); } - return offsetDer.reshape(mean.shape as ShapeMap[R]); + return reshape(offsetDer, mean.shape as ShapeMap[R]); }; return { x: derX, diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts index f66124cc33e..2c680a2c387 100644 --- a/tfjs-core/src/kernel_names.ts +++ b/tfjs-core/src/kernel_names.ts @@ -21,10 +21,10 @@ import {NamedTensorInfoMap} from './kernel_registry'; import {PixelData} from './types'; -export const BatchNormalization = 'BatchNormalization'; -export type BatchNormalizationInputs = +export const FusedBatchNorm = 'FusedBatchNorm'; +export type FusedBatchNormInputs = Pick; -export interface BatchNormalizationAttrs { +export interface FusedBatchNormAttrs { varianceEpsilon: number; } diff --git a/tfjs-core/src/ops/batchnorm.ts b/tfjs-core/src/ops/batchnorm.ts index d2475908d5e..d707356ceaf 100644 --- a/tfjs-core/src/ops/batchnorm.ts +++ b/tfjs-core/src/ops/batchnorm.ts @@ -16,7 +16,7 @@ */ import {ENGINE, ForwardFunc} from '../engine'; -import {BatchNormalizationAttrs, BatchNormalizationInputs} from '../kernel_names'; +import {FusedBatchNormAttrs, FusedBatchNormInputs} from '../kernel_names'; import {NamedAttrMap} from '../kernel_registry'; import {Tensor, Tensor1D, Tensor4D} from '../tensor'; import {NamedTensorMap} from '../tensor_types'; @@ -99,24 +99,26 @@ function batchNorm_( () => 'Batch normalization gradient requires mean and scale to have ' + 'equal ranks.'); - const x4D: Tensor4D = xAs4D($x); - const forward: ForwardFunc = (backend, save) => { + const x4D: Tensor4D = xAs4D($x); + const res = backend.batchNormalization( x4D, as1DOr4D($mean), as1DOr4D($variance), varianceEpsilon, as1DOr4D($scale), as1DOr4D($offset)); + save([$x, $mean, $variance, $scale]); + return res; }; - const inputs: BatchNormalizationInputs = + const inputs: FusedBatchNormInputs = {x: $x, scale: $scale, offset: $offset, mean: $mean, variance: $variance}; - const attrs: BatchNormalizationAttrs = {varianceEpsilon}; + const attrs: FusedBatchNormAttrs = {varianceEpsilon}; const res = ENGINE.runKernelFunc( forward, inputs as {} as NamedTensorMap, null /* gradient */, - 'BatchNormalization', attrs as {} as NamedAttrMap); + 'FusedBatchNorm', attrs as {} as NamedAttrMap); return res.reshape($x.shape); } @@ -137,5 +139,6 @@ function as1DOr4D(x: Tensor): Tensor4D|Tensor1D { return x as Tensor4D; } +// todo(yassogba): Remove batchNormalization since it is deprecated. export const batchNormalization = op({batchNormalization_}); export const batchNorm = op({batchNorm_}); diff --git a/tfjs-core/src/ops/batchnorm2d.ts b/tfjs-core/src/ops/batchnorm2d.ts index 09711430d01..4f80363307a 100644 --- a/tfjs-core/src/ops/batchnorm2d.ts +++ b/tfjs-core/src/ops/batchnorm2d.ts @@ -19,9 +19,9 @@ import {convertToTensor} from '../tensor_util_env'; import {TensorLike} from '../types'; import * as util from '../util'; -import {op} from './operation'; import {batchNorm} from './batchnorm'; import {warnDeprecation} from './batchnorm_util'; +import {op} from './operation'; /** * Batch normalization, strictly for 2D. For the more relaxed version, see @@ -35,45 +35,44 @@ import {warnDeprecation} from './batchnorm_util'; * @param varianceEpsilon A small float number to avoid dividing by 0. */ function batchNorm2d_( - x: Tensor2D | TensorLike, mean: Tensor2D | Tensor1D | TensorLike, - variance: Tensor2D | Tensor1D | TensorLike, - offset?: Tensor2D | Tensor1D | TensorLike, - scale?: Tensor2D | Tensor1D | TensorLike, - varianceEpsilon?: number): Tensor2D { + x: Tensor2D|TensorLike, mean: Tensor2D|Tensor1D|TensorLike, + variance: Tensor2D|Tensor1D|TensorLike, + offset?: Tensor2D|Tensor1D|TensorLike, scale?: Tensor2D|Tensor1D|TensorLike, + varianceEpsilon?: number): Tensor2D { const $x = convertToTensor(x, 'x', 'batchNorm'); const $mean = convertToTensor(mean, 'mean', 'batchNorm'); const $variance = convertToTensor(variance, 'variance', 'batchNorm'); - let $scale: Tensor2D | Tensor1D; + let $scale: Tensor2D|Tensor1D; if (scale != null) { $scale = convertToTensor(scale, 'scale', 'batchNorm'); } - let $offset: Tensor2D | Tensor1D; + let $offset: Tensor2D|Tensor1D; if (offset != null) { $offset = convertToTensor(offset, 'offset', 'batchNorm'); } util.assert( - $x.rank === 2, - () => `Error in batchNorm3D: x must be rank 3 but got rank ` + - `${$x.rank}.`); + $x.rank === 2, + () => `Error in batchNorm3D: x must be rank 3 but got rank ` + + `${$x.rank}.`); util.assert( - $mean.rank === 2 || $mean.rank === 1, - () => `Error in batchNorm2D: mean must be rank 2 or rank 1 but ` + - `got rank ${$mean.rank}.`); + $mean.rank === 2 || $mean.rank === 1, + () => `Error in batchNorm2D: mean must be rank 2 or rank 1 but ` + + `got rank ${$mean.rank}.`); util.assert( - $variance.rank === 2 || $variance.rank === 1, - () => `Error in batchNorm2D: variance must be rank 2 or rank 1 ` + - `but got rank ${$variance.rank}.`); + $variance.rank === 2 || $variance.rank === 1, + () => `Error in batchNorm2D: variance must be rank 2 or rank 1 ` + + `but got rank ${$variance.rank}.`); if ($scale != null) { util.assert( - $scale.rank === 2 || $scale.rank === 1, - () => `Error in batchNorm2D: scale must be rank 2 or rank 1 ` + - `but got rank ${$scale.rank}.`); + $scale.rank === 2 || $scale.rank === 1, + () => `Error in batchNorm2D: scale must be rank 2 or rank 1 ` + + `but got rank ${$scale.rank}.`); } if ($offset != null) { util.assert( - $offset.rank === 2 || $offset.rank === 1, - () => `Error in batchNorm2D: offset must be rank 2 or rank 1 ` + - `but got rank ${$offset.rank}.`); + $offset.rank === 2 || $offset.rank === 1, + () => `Error in batchNorm2D: offset must be rank 2 or rank 1 ` + + `but got rank ${$offset.rank}.`); } return batchNorm($x, $mean, $variance, $offset, $scale, varianceEpsilon); @@ -84,13 +83,14 @@ function batchNorm2d_( * argument change of scale, offset, and varianceEpsilon. */ function batchNormalization2d_( - x: Tensor2D | TensorLike, mean: Tensor2D | Tensor1D | TensorLike, - variance: Tensor2D | Tensor1D | TensorLike, varianceEpsilon = .001, - scale?: Tensor2D | Tensor1D | TensorLike, - offset?: Tensor2D | Tensor1D | TensorLike): Tensor2D { + x: Tensor2D|TensorLike, mean: Tensor2D|Tensor1D|TensorLike, + variance: Tensor2D|Tensor1D|TensorLike, varianceEpsilon = .001, + scale?: Tensor2D|Tensor1D|TensorLike, + offset?: Tensor2D|Tensor1D|TensorLike): Tensor2D { warnDeprecation(); return batchNorm2d_(x, mean, variance, offset, scale, varianceEpsilon); } +// todo(yassogba): Remove batchNormalization2d since it is deprecated. export const batchNormalization2d = op({batchNormalization2d_}); export const batchNorm2d = op({batchNorm2d_}); diff --git a/tfjs-core/src/ops/batchnorm3d.ts b/tfjs-core/src/ops/batchnorm3d.ts index e0edfbf9e93..f8cf991e59d 100644 --- a/tfjs-core/src/ops/batchnorm3d.ts +++ b/tfjs-core/src/ops/batchnorm3d.ts @@ -19,9 +19,10 @@ import {convertToTensor} from '../tensor_util_env'; import {TensorLike} from '../types'; import * as util from '../util'; -import {op} from './operation'; import {batchNorm} from './batchnorm'; import {warnDeprecation} from './batchnorm_util'; +import {op} from './operation'; + /** * Batch normalization, strictly for 3D. For the more relaxed version, see * `tf.batchNorm`. @@ -34,45 +35,44 @@ import {warnDeprecation} from './batchnorm_util'; * @param varianceEpsilon A small float number to avoid dividing by 0. */ function batchNorm3d_( - x: Tensor3D | TensorLike, mean: Tensor3D | Tensor1D | TensorLike, - variance: Tensor3D | Tensor1D | TensorLike, - offset?: Tensor3D | Tensor1D | TensorLike, - scale?: Tensor3D | Tensor1D | TensorLike, - varianceEpsilon?: number): Tensor3D { + x: Tensor3D|TensorLike, mean: Tensor3D|Tensor1D|TensorLike, + variance: Tensor3D|Tensor1D|TensorLike, + offset?: Tensor3D|Tensor1D|TensorLike, scale?: Tensor3D|Tensor1D|TensorLike, + varianceEpsilon?: number): Tensor3D { const $x = convertToTensor(x, 'x', 'batchNorm'); const $mean = convertToTensor(mean, 'mean', 'batchNorm'); const $variance = convertToTensor(variance, 'variance', 'batchNorm'); - let $scale: Tensor3D | Tensor1D; + let $scale: Tensor3D|Tensor1D; if (scale != null) { $scale = convertToTensor(scale, 'scale', 'batchNorm'); } - let $offset: Tensor3D | Tensor1D; + let $offset: Tensor3D|Tensor1D; if (offset != null) { $offset = convertToTensor(offset, 'offset', 'batchNorm'); } util.assert( - $x.rank === 3, - () => `Error in batchNorm3D: x must be rank 3 but got rank ` + - `${$x.rank}.`); + $x.rank === 3, + () => `Error in batchNorm3D: x must be rank 3 but got rank ` + + `${$x.rank}.`); util.assert( - $mean.rank === 3 || $mean.rank === 1, - () => `Error in batchNorm3D: mean must be rank 3 or rank 1 but ` + - `got rank ${$mean.rank}.`); + $mean.rank === 3 || $mean.rank === 1, + () => `Error in batchNorm3D: mean must be rank 3 or rank 1 but ` + + `got rank ${$mean.rank}.`); util.assert( - $variance.rank === 3 || $variance.rank === 1, - () => `Error in batchNorm3D: variance must be rank 3 or rank 1 ` + - `but got rank ${$variance.rank}.`); + $variance.rank === 3 || $variance.rank === 1, + () => `Error in batchNorm3D: variance must be rank 3 or rank 1 ` + + `but got rank ${$variance.rank}.`); if ($scale != null) { util.assert( - $scale.rank === 3 || $scale.rank === 1, - () => `Error in batchNorm3D: scale must be rank 3 or rank 1 ` + - `but got rank ${$scale.rank}.`); + $scale.rank === 3 || $scale.rank === 1, + () => `Error in batchNorm3D: scale must be rank 3 or rank 1 ` + + `but got rank ${$scale.rank}.`); } if ($offset != null) { util.assert( - $offset.rank === 3 || $offset.rank === 1, - () => `Error in batchNorm3D: offset must be rank 3 or rank 1 ` + - `but got rank ${$offset.rank}.`); + $offset.rank === 3 || $offset.rank === 1, + () => `Error in batchNorm3D: offset must be rank 3 or rank 1 ` + + `but got rank ${$offset.rank}.`); } return batchNorm($x, $mean, $variance, $offset, $scale, varianceEpsilon); @@ -83,13 +83,14 @@ function batchNorm3d_( * argument change of scale, offset, and varianceEpsilon. */ function batchNormalization3d_( - x: Tensor3D | TensorLike, mean: Tensor3D | Tensor1D | TensorLike, - variance: Tensor3D | Tensor1D | TensorLike, varianceEpsilon = .001, - scale?: Tensor3D | Tensor1D | TensorLike, - offset?: Tensor3D | Tensor1D | TensorLike): Tensor3D { + x: Tensor3D|TensorLike, mean: Tensor3D|Tensor1D|TensorLike, + variance: Tensor3D|Tensor1D|TensorLike, varianceEpsilon = .001, + scale?: Tensor3D|Tensor1D|TensorLike, + offset?: Tensor3D|Tensor1D|TensorLike): Tensor3D { warnDeprecation(); return batchNorm3d_(x, mean, variance, offset, scale, varianceEpsilon); } +// todo(yassogba): Remove batchNormalization3d since it is deprecated. export const batchNormalization3d = op({batchNormalization3d_}); export const batchNorm3d = op({batchNorm3d_}); diff --git a/tfjs-core/src/ops/batchnorm4d.ts b/tfjs-core/src/ops/batchnorm4d.ts index e46bfa34125..bbb56d0bb33 100644 --- a/tfjs-core/src/ops/batchnorm4d.ts +++ b/tfjs-core/src/ops/batchnorm4d.ts @@ -19,9 +19,10 @@ import {convertToTensor} from '../tensor_util_env'; import {TensorLike} from '../types'; import * as util from '../util'; -import {op} from './operation'; import {batchNorm} from './batchnorm'; import {warnDeprecation} from './batchnorm_util'; +import {op} from './operation'; + /** * Batch normalization, strictly for 4D. For the more relaxed version, see * `tf.batchNorm`. @@ -34,45 +35,44 @@ import {warnDeprecation} from './batchnorm_util'; * @param varianceEpsilon A small float number to avoid dividing by 0. */ function batchNorm4d_( - x: Tensor4D | TensorLike, mean: Tensor4D | Tensor1D | TensorLike, - variance: Tensor4D | Tensor1D | TensorLike, - offset?: Tensor4D | Tensor1D | TensorLike, - scale?: Tensor4D | Tensor1D | TensorLike, - varianceEpsilon?: number): Tensor4D { + x: Tensor4D|TensorLike, mean: Tensor4D|Tensor1D|TensorLike, + variance: Tensor4D|Tensor1D|TensorLike, + offset?: Tensor4D|Tensor1D|TensorLike, scale?: Tensor4D|Tensor1D|TensorLike, + varianceEpsilon?: number): Tensor4D { const $x = convertToTensor(x, 'x', 'batchNorm'); const $mean = convertToTensor(mean, 'mean', 'batchNorm'); const $variance = convertToTensor(variance, 'variance', 'batchNorm'); - let $scale: Tensor4D | Tensor1D; + let $scale: Tensor4D|Tensor1D; if (scale != null) { $scale = convertToTensor(scale, 'scale', 'batchNorm'); } - let $offset: Tensor4D | Tensor1D; + let $offset: Tensor4D|Tensor1D; if (offset != null) { $offset = convertToTensor(offset, 'offset', 'batchNorm'); } util.assert( - $x.rank === 4, - () => `Error in batchNorm4D: x must be rank 4 but got rank ` + - `${$x.rank}.`); + $x.rank === 4, + () => `Error in batchNorm4D: x must be rank 4 but got rank ` + + `${$x.rank}.`); util.assert( - $mean.rank === 4 || $mean.rank === 1, - () => `Error in batchNorm4D: mean must be rank 4 or rank 1 but ` + - `got rank ${$mean.rank}.`); + $mean.rank === 4 || $mean.rank === 1, + () => `Error in batchNorm4D: mean must be rank 4 or rank 1 but ` + + `got rank ${$mean.rank}.`); util.assert( - $variance.rank === 4 || $variance.rank === 1, - () => `Error in batchNorm4D: variance must be rank 4 or rank 1 ` + - `but got rank ${$variance.rank}.`); + $variance.rank === 4 || $variance.rank === 1, + () => `Error in batchNorm4D: variance must be rank 4 or rank 1 ` + + `but got rank ${$variance.rank}.`); if ($scale != null) { util.assert( - $scale.rank === 4 || $scale.rank === 1, - () => `Error in batchNorm4D: scale must be rank 4 or rank 1 ` + - `but got rank ${$scale.rank}.`); + $scale.rank === 4 || $scale.rank === 1, + () => `Error in batchNorm4D: scale must be rank 4 or rank 1 ` + + `but got rank ${$scale.rank}.`); } if ($offset != null) { util.assert( - $offset.rank === 4 || $offset.rank === 1, - () => `Error in batchNorm4D: offset must be rank 4 or rank 1 ` + - `but got rank ${$offset.rank}.`); + $offset.rank === 4 || $offset.rank === 1, + () => `Error in batchNorm4D: offset must be rank 4 or rank 1 ` + + `but got rank ${$offset.rank}.`); } return batchNorm($x, $mean, $variance, $offset, $scale, varianceEpsilon); } @@ -82,13 +82,14 @@ function batchNorm4d_( * argument change of scale, offset, and varianceEpsilon. */ function batchNormalization4d_( - x: Tensor4D | TensorLike, mean: Tensor4D | Tensor1D | TensorLike, - variance: Tensor4D | Tensor1D | TensorLike, varianceEpsilon = .001, - scale?: Tensor4D | Tensor1D | TensorLike, - offset?: Tensor4D | Tensor1D | TensorLike): Tensor4D { + x: Tensor4D|TensorLike, mean: Tensor4D|Tensor1D|TensorLike, + variance: Tensor4D|Tensor1D|TensorLike, varianceEpsilon = .001, + scale?: Tensor4D|Tensor1D|TensorLike, + offset?: Tensor4D|Tensor1D|TensorLike): Tensor4D { warnDeprecation(); return batchNorm4d_(x, mean, variance, offset, scale, varianceEpsilon); } +// todo(yassogba): Remove batchNormalization4d since it is deprecated. export const batchNormalization4d = op({batchNormalization4d_}); export const batchNorm4d = op({batchNorm4d_}); diff --git a/tfjs-core/src/register_all_gradients.ts b/tfjs-core/src/register_all_gradients.ts index 8db20b82768..35974b826ac 100644 --- a/tfjs-core/src/register_all_gradients.ts +++ b/tfjs-core/src/register_all_gradients.ts @@ -14,9 +14,9 @@ * limitations under the License. * ============================================================================= */ -import {batchNormalizationGradConfig} from './gradients/BatchNormalization_grad'; import {broadcastToGradConfig} from './gradients/BroadcastTo_grad'; import {divGradConfig} from './gradients/Div_grad'; +import {fusedBatchNormGradConfig} from './gradients/FusedBatchNorm_grad'; import {identityGradConfig} from './gradients/Identity_grad'; import {oneHotGradConfig} from './gradients/OneHot_grad'; import {padV2GradConfig} from './gradients/PadV2_grad'; @@ -29,9 +29,9 @@ import {registerGradient} from './kernel_registry'; // Export all kernel configs here so that the package can auto register them const gradConfigs: GradConfig[] = [ - batchNormalizationGradConfig, divGradConfig, squareGradConfig, - squaredDifferenceGradConfig, broadcastToGradConfig, identityGradConfig, - tileGradConfig, oneHotGradConfig, transposeGradConfig, padV2GradConfig + broadcastToGradConfig, divGradConfig, fusedBatchNormGradConfig, + identityGradConfig, oneHotGradConfig, padV2GradConfig, squareGradConfig, + squaredDifferenceGradConfig, tileGradConfig, transposeGradConfig ]; for (const gradientConfig of gradConfigs) { From 8b53a13ca43a463dc311560232de6c3b16325cf7 Mon Sep 17 00:00:00 2001 From: Na Li Date: Tue, 31 Mar 2020 12:25:19 -0700 Subject: [PATCH 8/8] Change to use FusedBatchNorm as kernel name for wasm. --- .../src/kernels/FusedBatchNorm.ts | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tfjs-backend-wasm/src/kernels/FusedBatchNorm.ts b/tfjs-backend-wasm/src/kernels/FusedBatchNorm.ts index 03c82decb15..ac68da3bee7 100644 --- a/tfjs-backend-wasm/src/kernels/FusedBatchNorm.ts +++ b/tfjs-backend-wasm/src/kernels/FusedBatchNorm.ts @@ -32,19 +32,19 @@ interface BatchNormAttrs extends NamedAttrMap { } let wasmBatchNorm: ( - xId: number, meanId: number, varianceId: number, offsetId: number, - scaleId: number, varianceEpsilon: number, outId: number) => void; + xId: number, meanId: number, varianceId: number, offsetId: number, + scaleId: number, varianceEpsilon: number, outId: number) => void; function setup(backend: BackendWasm): void { wasmBatchNorm = backend.wasm.cwrap( - 'FusedBatchNorm', null /* void */, - ['number', 'number', 'number', 'number', 'number', 'number', 'number']); + 'FusedBatchNorm', null /* void */, + ['number', 'number', 'number', 'number', 'number', 'number', 'number']); } function fusedBatchNorm( - args: - {backend: BackendWasm, inputs: BatchNormInputs, attrs: BatchNormAttrs}): - TensorInfo { + args: + {backend: BackendWasm, inputs: BatchNormInputs, attrs: BatchNormAttrs}): + TensorInfo { const {backend, inputs, attrs} = args; const {varianceEpsilon} = attrs; const {x, mean, variance, offset, scale} = inputs; @@ -63,12 +63,12 @@ function fusedBatchNorm( const outId = backend.dataIdMap.get(out.dataId).id; wasmBatchNorm( - xId, meanId, varianceId, offsetId, scaleId, varianceEpsilon, outId); + xId, meanId, varianceId, offsetId, scaleId, varianceEpsilon, outId); return out; } registerKernel({ - kernelName: 'BatchNormalization', + kernelName: 'FusedBatchNorm', backendName: 'wasm', setupFunc: setup, kernelFunc: fusedBatchNorm