@tensorflow/tfjs-node
Version:
This repository provides native TensorFlow execution in backend JavaScript applications under the Node.js runtime, accelerated by the TensorFlow C binary under the hood. It provides the same API as [TensorFlow.js](https://js.tensorflow.org/api/latest/).
67 lines (66 loc) • 2.89 kB
JavaScript
;
/**
* @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.
* =============================================================================
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.fusedBatchNormConfig = void 0;
var tfjs_1 = require("@tensorflow/tfjs");
var nodejs_kernel_backend_1 = require("../nodejs_kernel_backend");
exports.fusedBatchNormConfig = {
kernelName: tfjs_1.FusedBatchNorm,
backendName: 'tensorflow',
kernelFunc: function (args) {
var _a = args.inputs, x = _a.x, mean = _a.mean, variance = _a.variance;
var _b = args.inputs, scale = _b.scale, offset = _b.offset;
var backend = args.backend;
var varianceEpsilon = args.attrs.varianceEpsilon;
return (0, tfjs_1.tidy)(function () {
if (mean.rank > 1) {
// Fused batch norm doesn't work with high-dim mean/var/scale/offset.
var inv = (0, tfjs_1.rsqrt)((0, tfjs_1.add)(variance, (0, tfjs_1.scalar)(varianceEpsilon)));
if (scale != null) {
inv = (0, tfjs_1.mul)(inv, scale);
}
var xNorm = (0, tfjs_1.mul)((0, tfjs_1.sub)(x, mean), inv);
return offset != null ? (0, tfjs_1.add)(xNorm, offset) : xNorm;
}
var dataFormat = 'NHWC';
var depth = x.shape[3];
var opAttrs = [
(0, nodejs_kernel_backend_1.createTensorsTypeOpAttr)('T', x.dtype),
{
name: 'epsilon',
type: backend.binding.TF_ATTR_FLOAT,
value: varianceEpsilon
},
{
name: 'data_format',
type: backend.binding.TF_ATTR_STRING,
value: dataFormat
},
{ name: 'is_training', type: backend.binding.TF_ATTR_BOOL, value: false },
];
var numOutputs = 5;
if (scale == null) {
scale = (0, tfjs_1.fill)([depth], 1);
}
if (offset == null) {
offset = (0, tfjs_1.fill)([depth], 0);
}
return backend.executeMultipleOutputs(tfjs_1.FusedBatchNorm, opAttrs, [x, scale, offset, mean, variance], numOutputs)[0];
});
}
};