@tensorflow/tfjs-backend-wasm
Version:
This package adds a WebAssembly backend to TensorFlow.js. It currently supports the following models from our [models](https://github.com/tensorflow/tfjs-models) repo: - BlazeFace - BodyPix - CocoSSD - Face landmarks detection - HandPose - KNN classifier
78 lines • 10.9 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.
* =============================================================================
*/
import { backend_util, Mean, util } from '@tensorflow/tfjs-core';
import { cast } from './Cast';
import { permuteAxesAndTranspose } from './kernel_utils';
let wasmMean;
function setup(backend) {
wasmMean =
backend.wasm.cwrap(Mean, null /*void*/, ['number, number, number']);
}
export function mean(args) {
const { backend, inputs, attrs } = args;
const { axis, keepDims } = attrs;
const { x } = inputs;
const xId = backend.dataIdMap.get(x.dataId).id;
let inputId = xId;
let input = x;
const { transposed, axes, originalAxes, inputWasTransposed } = permuteAxesAndTranspose(x, axis, backend);
let reductionAxes = axes;
if (inputWasTransposed) {
const transposedId = backend.dataIdMap.get(transposed.dataId).id;
if (transposedId !== xId) {
// transpose was not a no-op. We will need to dispose of this
// once we are done.
input = transposed;
inputId = transposedId;
reductionAxes = backend_util.getInnerMostAxes(reductionAxes.length, input.shape.length);
}
}
backend_util.assertAxesAreInnerMostDims('mean', reductionAxes, input.shape.length);
const [outShape, reduceShape] = backend_util.computeOutAndReduceShapes(input.shape, reductionAxes);
const reduceSize = util.sizeFromShape(reduceShape);
let castedInput = input;
if (input.dtype !== 'float32') {
castedInput =
cast({ backend, inputs: { x: input }, attrs: { dtype: 'float32' } });
inputId = backend.dataIdMap.get(castedInput.dataId).id;
}
const out = backend.makeOutput(outShape, 'float32');
if (util.sizeFromShape(input.shape) !== 0) {
const outId = backend.dataIdMap.get(out.dataId).id;
wasmMean(inputId, reduceSize, outId);
}
if (inputWasTransposed) {
// dispose of the transposed tensor.
backend.disposeData(transposed.dataId);
}
if (keepDims) {
// reshape
const newShape = backend_util.expandShapeToKeepDim(out.shape, originalAxes);
out.shape = newShape;
}
if (input.dtype !== 'float32') {
backend.disposeData(castedInput.dataId);
}
return out;
}
export const meanConfig = {
kernelName: Mean,
backendName: 'wasm',
setupFunc: setup,
kernelFunc: mean
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiTWVhbi5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uLy4uLy4uLy4uL3RmanMtYmFja2VuZC13YXNtL3NyYy9rZXJuZWxzL01lYW4udHMiXSwibmFtZXMiOltdLCJtYXBwaW5ncyI6IkFBQUE7Ozs7Ozs7Ozs7Ozs7OztHQWVHO0FBRUgsT0FBTyxFQUFDLFlBQVksRUFBNEIsSUFBSSxFQUFxQyxJQUFJLEVBQUMsTUFBTSx1QkFBdUIsQ0FBQztBQUc1SCxPQUFPLEVBQUMsSUFBSSxFQUFDLE1BQU0sUUFBUSxDQUFDO0FBRTVCLE9BQU8sRUFBQyx1QkFBdUIsRUFBQyxNQUFNLGdCQUFnQixDQUFDO0FBRXZELElBQUksUUFBa0UsQ0FBQztBQUV2RSxTQUFTLEtBQUssQ0FBQyxPQUFvQjtJQUNqQyxRQUFRO1FBQ0osT0FBTyxDQUFDLElBQUksQ0FBQyxLQUFLLENBQUMsSUFBSSxFQUFFLElBQUksQ0FBQyxRQUFRLEVBQUUsQ0FBQyx3QkFBd0IsQ0FBQyxDQUFDLENBQUM7QUFDMUUsQ0FBQztBQUVELE1BQU0sVUFBVSxJQUFJLENBQ2hCLElBQWtFO0lBRXBFLE1BQU0sRUFBQyxPQUFPLEVBQUUsTUFBTSxFQUFFLEtBQUssRUFBQyxHQUFHLElBQUksQ0FBQztJQUN0QyxNQUFNLEVBQUMsSUFBSSxFQUFFLFFBQVEsRUFBQyxHQUFHLEtBQUssQ0FBQztJQUMvQixNQUFNLEVBQUMsQ0FBQyxFQUFDLEdBQUcsTUFBTSxDQUFDO0lBQ25CLE1BQU0sR0FBRyxHQUFHLE9BQU8sQ0FBQyxTQUFTLENBQUMsR0FBRyxDQUFDLENBQUMsQ0FBQyxNQUFNLENBQUMsQ0FBQyxFQUFFLENBQUM7SUFDL0MsSUFBSSxPQUFPLEdBQUcsR0FBRyxDQUFDO0lBQ2xCLElBQUksS0FBSyxHQUFHLENBQUMsQ0FBQztJQUVkLE1BQU0sRUFBQyxVQUFVLEVBQUUsSUFBSSxFQUFFLFlBQVksRUFBRSxrQkFBa0IsRUFBQyxHQUN0RCx1QkFBdUIsQ0FBQyxDQUFDLEVBQUUsSUFBSSxFQUFFLE9BQU8sQ0FBQyxDQUFDO0lBRTlDLElBQUksYUFBYSxHQUFHLElBQUksQ0FBQztJQUN6QixJQUFJLGtCQUFrQixFQUFFO1FBQ3RCLE1BQU0sWUFBWSxHQUFHLE9BQU8sQ0FBQyxTQUFTLENBQUMsR0FBRyxDQUFDLFVBQVUsQ0FBQyxNQUFNLENBQUMsQ0FBQyxFQUFFLENBQUM7UUFDakUsSUFBSSxZQUFZLEtBQUssR0FBRyxFQUFFO1lBQ3hCLDZEQUE2RDtZQUM3RCxvQkFBb0I7WUFDcEIsS0FBSyxHQUFHLFVBQVUsQ0FBQztZQUNuQixPQUFPLEdBQUcsWUFBWSxDQUFDO1lBQ3ZCLGFBQWEsR0FBRyxZQUFZLENBQUMsZ0JBQWdCLENBQ3pDLGFBQWEsQ0FBQyxNQUFNLEVBQUUsS0FBSyxDQUFDLEtBQUssQ0FBQyxNQUFNLENBQUMsQ0FBQztTQUMvQztLQUNGO0lBRUQsWUFBWSxDQUFDLDBCQUEwQixDQUNuQyxNQUFNLEVBQUUsYUFBYSxFQUFFLEtBQUssQ0FBQyxLQUFLLENBQUMsTUFBTSxDQUFDLENBQUM7SUFDL0MsTUFBTSxDQUFDLFFBQVEsRUFBRSxXQUFXLENBQUMsR0FDekIsWUFBWSxDQUFDLHlCQUF5QixDQUFDLEtBQUssQ0FBQyxLQUFLLEVBQUUsYUFBYSxDQUFDLENBQUM7SUFDdkUsTUFBTSxVQUFVLEdBQUcsSUFBSSxDQUFDLGFBQWEsQ0FBQyxXQUFXLENBQUMsQ0FBQztJQUNuRCxJQUFJLFdBQVcsR0FBRyxLQUFLLENBQUM7SUFDeEIsSUFBSSxLQUFLLENBQUMsS0FBSyxLQUFLLFNBQVMsRUFBRTtRQUM3QixXQUFXO1lBQ1AsSUFBSSxDQUFDLEVBQUMsT0FBTyxFQUFFLE1BQU0sRUFBRSxFQUFDLENBQUMsRUFBRSxLQUFLLEVBQUMsRUFBRSxLQUFLLEVBQUUsRUFBQyxLQUFLLEVBQUUsU0FBUyxFQUFDLEVBQUMsQ0FBQyxDQUFDO1FBQ25FLE9BQU8sR0FBRyxPQUFPLENBQUMsU0FBUyxDQUFDLEdBQUcsQ0FBQyxXQUFXLENBQUMsTUFBTSxDQUFDLENBQUMsRUFBRSxDQUFDO0tBQ3hEO0lBRUQsTUFBTSxHQUFHLEdBQUcsT0FBTyxDQUFDLFVBQVUsQ0FBQyxRQUFRLEVBQUUsU0FBUyxDQUFDLENBQUM7SUFDcEQsSUFBSSxJQUFJLENBQUMsYUFBYSxDQUFDLEtBQUssQ0FBQyxLQUFLLENBQUMsS0FBSyxDQUFDLEVBQUU7UUFDekMsTUFBTSxLQUFLLEdBQUcsT0FBTyxDQUFDLFNBQVMsQ0FBQyxHQUFHLENBQUMsR0FBRyxDQUFDLE1BQU0sQ0FBQyxDQUFDLEVBQUUsQ0FBQztRQUNuRCxRQUFRLENBQUMsT0FBTyxFQUFFLFVBQVUsRUFBRSxLQUFLLENBQUMsQ0FBQztLQUN0QztJQUVELElBQUksa0JBQWtCLEVBQUU7UUFDdEIsb0NBQW9DO1FBQ3BDLE9BQU8sQ0FBQyxXQUFXLENBQUMsVUFBVSxDQUFDLE1BQU0sQ0FBQyxDQUFDO0tBQ3hDO0lBRUQsSUFBSSxRQUFRLEVBQUU7UUFDWixVQUFVO1FBQ1YsTUFBTSxRQUFRLEdBQUcsWUFBWSxDQUFDLG9CQUFvQixDQUFDLEdBQUcsQ0FBQyxLQUFLLEVBQUUsWUFBWSxDQUFDLENBQUM7UUFDNUUsR0FBRyxDQUFDLEtBQUssR0FBRyxRQUFRLENBQUM7S0FDdEI7SUFFRCxJQUFJLEtBQUssQ0FBQyxLQUFLLEtBQUssU0FBUyxFQUFFO1FBQzdCLE9BQU8sQ0FBQyxXQUFXLENBQUMsV0FBVyxDQUFDLE1BQU0sQ0FBQyxDQUFDO0tBQ3pDO0lBRUQsT0FBTyxHQUFHLENBQUM7QUFDYixDQUFDO0FBRUQsTUFBTSxDQUFDLE1BQU0sVUFBVSxHQUFpQjtJQUN0QyxVQUFVLEVBQUUsSUFBSTtJQUNoQixXQUFXLEVBQUUsTUFBTTtJQUNuQixTQUFTLEVBQUUsS0FBSztJQUNoQixVQUFVLEVBQUUsSUFBNkI7Q0FDMUMsQ0FBQyIsInNvdXJjZXNDb250ZW50IjpbIi8qKlxuICogQGxpY2Vuc2VcbiAqIENvcHlyaWdodCAyMDIwIEdvb2dsZSBMTEMuIEFsbCBSaWdodHMgUmVzZXJ2ZWQuXG4gKiBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgXCJMaWNlbnNlXCIpO1xuICogeW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLlxuICogWW91IG1heSBvYnRhaW4gYSBjb3B5IG9mIHRoZSBMaWNlbnNlIGF0XG4gKlxuICogaHR0cDovL3d3dy5hcGFjaGUub3JnL2xpY2Vuc2VzL0xJQ0VOU0UtMi4wXG4gKlxuICogVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZVxuICogZGlzdHJpYnV0ZWQgdW5kZXIgdGhlIExpY2Vuc2UgaXMgZGlzdHJpYnV0ZWQgb24gYW4gXCJBUyBJU1wiIEJBU0lTLFxuICogV0lUSE9VVCBXQVJSQU5USUVTIE9SIENPTkRJVElPTlMgT0YgQU5ZIEtJTkQsIGVpdGhlciBleHByZXNzIG9yIGltcGxpZWQuXG4gKiBTZWUgdGhlIExpY2Vuc2UgZm9yIHRoZSBzcGVjaWZpYyBsYW5ndWFnZSBnb3Zlcm5pbmcgcGVybWlzc2lvbnMgYW5kXG4gKiBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS5cbiAqID09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09XG4gKi9cblxuaW1wb3J0IHtiYWNrZW5kX3V0aWwsIEtlcm5lbENvbmZpZywgS2VybmVsRnVuYywgTWVhbiwgTWVhbkF0dHJzLCBNZWFuSW5wdXRzLCBUZW5zb3JJbmZvLCB1dGlsfSBmcm9tICdAdGVuc29yZmxvdy90ZmpzLWNvcmUnO1xuXG5pbXBvcnQge0JhY2tlbmRXYXNtfSBmcm9tICcuLi9iYWNrZW5kX3dhc20nO1xuaW1wb3J0IHtjYXN0fSBmcm9tICcuL0Nhc3QnO1xuXG5pbXBvcnQge3Blcm11dGVBeGVzQW5kVHJhbnNwb3NlfSBmcm9tICcuL2tlcm5lbF91dGlscyc7XG5cbmxldCB3YXNtTWVhbjogKHhJZDogbnVtYmVyLCByZWR1Y2VTaXplOiBudW1iZXIsIG91dElkOiBudW1iZXIpID0+IHZvaWQ7XG5cbmZ1bmN0aW9uIHNldHVwKGJhY2tlbmQ6IEJhY2tlbmRXYXNtKTogdm9pZCB7XG4gIHdhc21NZWFuID1cbiAgICAgIGJhY2tlbmQud2FzbS5jd3JhcChNZWFuLCBudWxsIC8qdm9pZCovLCBbJ251bWJlciwgbnVtYmVyLCBudW1iZXInXSk7XG59XG5cbmV4cG9ydCBmdW5jdGlvbiBtZWFuKFxuICAgIGFyZ3M6IHtiYWNrZW5kOiBCYWNrZW5kV2FzbSwgaW5wdXRzOiBNZWFuSW5wdXRzLCBhdHRyczogTWVhbkF0dHJzfSk6XG4gICAgVGVuc29ySW5mbyB7XG4gIGNvbnN0IHtiYWNrZW5kLCBpbnB1dHMsIGF0dHJzfSA9IGFyZ3M7XG4gIGNvbnN0IHtheGlzLCBrZWVwRGltc30gPSBhdHRycztcbiAgY29uc3Qge3h9ID0gaW5wdXRzO1xuICBjb25zdCB4SWQgPSBiYWNrZW5kLmRhdGFJZE1hcC5nZXQoeC5kYXRhSWQpLmlkO1xuICBsZXQgaW5wdXRJZCA9IHhJZDtcbiAgbGV0IGlucHV0ID0geDtcblxuICBjb25zdCB7dHJhbnNwb3NlZCwgYXhlcywgb3JpZ2luYWxBeGVzLCBpbnB1dFdhc1RyYW5zcG9zZWR9ID1cbiAgICAgIHBlcm11dGVBeGVzQW5kVHJhbnNwb3NlKHgsIGF4aXMsIGJhY2tlbmQpO1xuXG4gIGxldCByZWR1Y3Rpb25BeGVzID0gYXhlcztcbiAgaWYgKGlucHV0V2FzVHJhbnNwb3NlZCkge1xuICAgIGNvbnN0IHRyYW5zcG9zZWRJZCA9IGJhY2tlbmQuZGF0YUlkTWFwLmdldCh0cmFuc3Bvc2VkLmRhdGFJZCkuaWQ7XG4gICAgaWYgKHRyYW5zcG9zZWRJZCAhPT0geElkKSB7XG4gICAgICAvLyB0cmFuc3Bvc2Ugd2FzIG5vdCBhIG5vLW9wLiBXZSB3aWxsIG5lZWQgdG8gZGlzcG9zZSBvZiB0aGlzXG4gICAgICAvLyBvbmNlIHdlIGFyZSBkb25lLlxuICAgICAgaW5wdXQgPSB0cmFuc3Bvc2VkO1xuICAgICAgaW5wdXRJZCA9IHRyYW5zcG9zZWRJZDtcbiAgICAgIHJlZHVjdGlvbkF4ZXMgPSBiYWNrZW5kX3V0aWwuZ2V0SW5uZXJNb3N0QXhlcyhcbiAgICAgICAgICByZWR1Y3Rpb25BeGVzLmxlbmd0aCwgaW5wdXQuc2hhcGUubGVuZ3RoKTtcbiAgICB9XG4gIH1cblxuICBiYWNrZW5kX3V0aWwuYXNzZXJ0QXhlc0FyZUlubmVyTW9zdERpbXMoXG4gICAgICAnbWVhbicsIHJlZHVjdGlvbkF4ZXMsIGlucHV0LnNoYXBlLmxlbmd0aCk7XG4gIGNvbnN0IFtvdXRTaGFwZSwgcmVkdWNlU2hhcGVdID1cbiAgICAgIGJhY2tlbmRfdXRpbC5jb21wdXRlT3V0QW5kUmVkdWNlU2hhcGVzKGlucHV0LnNoYXBlLCByZWR1Y3Rpb25BeGVzKTtcbiAgY29uc3QgcmVkdWNlU2l6ZSA9IHV0aWwuc2l6ZUZyb21TaGFwZShyZWR1Y2VTaGFwZSk7XG4gIGxldCBjYXN0ZWRJbnB1dCA9IGlucHV0O1xuICBpZiAoaW5wdXQuZHR5cGUgIT09ICdmbG9hdDMyJykge1xuICAgIGNhc3RlZElucHV0ID1cbiAgICAgICAgY2FzdCh7YmFja2VuZCwgaW5wdXRzOiB7eDogaW5wdXR9LCBhdHRyczoge2R0eXBlOiAnZmxvYXQzMid9fSk7XG4gICAgaW5wdXRJZCA9IGJhY2tlbmQuZGF0YUlkTWFwLmdldChjYXN0ZWRJbnB1dC5kYXRhSWQpLmlkO1xuICB9XG5cbiAgY29uc3Qgb3V0ID0gYmFja2VuZC5tYWtlT3V0cHV0KG91dFNoYXBlLCAnZmxvYXQzMicpO1xuICBpZiAodXRpbC5zaXplRnJvbVNoYXBlKGlucHV0LnNoYXBlKSAhPT0gMCkge1xuICAgIGNvbnN0IG91dElkID0gYmFja2VuZC5kYXRhSWRNYXAuZ2V0KG91dC5kYXRhSWQpLmlkO1xuICAgIHdhc21NZWFuKGlucHV0SWQsIHJlZHVjZVNpemUsIG91dElkKTtcbiAgfVxuXG4gIGlmIChpbnB1dFdhc1RyYW5zcG9zZWQpIHtcbiAgICAvLyBkaXNwb3NlIG9mIHRoZSB0cmFuc3Bvc2VkIHRlbnNvci5cbiAgICBiYWNrZW5kLmRpc3Bvc2VEYXRhKHRyYW5zcG9zZWQuZGF0YUlkKTtcbiAgfVxuXG4gIGlmIChrZWVwRGltcykge1xuICAgIC8vIHJlc2hhcGVcbiAgICBjb25zdCBuZXdTaGFwZSA9IGJhY2tlbmRfdXRpbC5leHBhbmRTaGFwZVRvS2VlcERpbShvdXQuc2hhcGUsIG9yaWdpbmFsQXhlcyk7XG4gICAgb3V0LnNoYXBlID0gbmV3U2hhcGU7XG4gIH1cblxuICBpZiAoaW5wdXQuZHR5cGUgIT09ICdmbG9hdDMyJykge1xuICAgIGJhY2tlbmQuZGlzcG9zZURhdGEoY2FzdGVkSW5wdXQuZGF0YUlkKTtcbiAgfVxuXG4gIHJldHVybiBvdXQ7XG59XG5cbmV4cG9ydCBjb25zdCBtZWFuQ29uZmlnOiBLZXJuZWxDb25maWcgPSB7XG4gIGtlcm5lbE5hbWU6IE1lYW4sXG4gIGJhY2tlbmROYW1lOiAnd2FzbScsXG4gIHNldHVwRnVuYzogc2V0dXAsXG4gIGtlcm5lbEZ1bmM6IG1lYW4gYXMgdW5rbm93biBhcyBLZXJuZWxGdW5jXG59O1xuIl19