@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
63 lines • 10.3 kB
JavaScript
/**
* @license
* Copyright 2019 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 { DepthToSpace, util } from '@tensorflow/tfjs-core';
let wasmDepthToSpace;
function setup(backend) {
wasmDepthToSpace = backend.wasm.cwrap(DepthToSpace, null /*void*/, [
'number',
'number',
'number',
'array',
'number',
'array',
'array',
'number',
'number', // outId
]);
}
export function depthToSpace(args) {
const { backend, inputs, attrs } = args;
const { x } = inputs;
const { blockSize, dataFormat } = attrs;
const batchSize = x.shape[0];
const inputHeight = (dataFormat === 'NHWC') ? x.shape[1] : x.shape[2];
const inputWidth = (dataFormat === 'NHWC') ? x.shape[2] : x.shape[3];
const inputDepth = (dataFormat === 'NHWC') ? x.shape[3] : x.shape[1];
const outputHeight = inputHeight * blockSize;
const outputWidth = inputWidth * blockSize;
const outputDepth = inputDepth / (blockSize * blockSize);
const outputShape = (dataFormat === 'NHWC') ?
[batchSize, outputHeight, outputWidth, outputDepth] :
[batchSize, outputDepth, outputHeight, outputWidth];
const out = backend.makeOutput(outputShape, 'float32');
const xData = backend.dataIdMap.get(x.dataId);
const xId = xData.id;
const xStridesBytes = new Uint8Array(new Int32Array(util.computeStrides(x.shape)).buffer);
const outputShapeBytes = new Uint8Array(new Int32Array(outputShape).buffer);
const outStridesBytes = new Uint8Array(new Int32Array(util.computeStrides(outputShape)).buffer);
const outId = backend.dataIdMap.get(out.dataId).id;
const channelsLast = dataFormat === 'NHWC' ? 1 : 0;
wasmDepthToSpace(xId, blockSize, channelsLast, xStridesBytes, x.shape.length - 1, outputShapeBytes, outStridesBytes, outputShape.length, outId);
return out;
}
export const depthToSpaceConfig = {
kernelName: DepthToSpace,
backendName: 'wasm',
setupFunc: setup,
kernelFunc: depthToSpace
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiRGVwdGhUb1NwYWNlLmpzIiwic291cmNlUm9vdCI6IiIsInNvdXJjZXMiOlsiLi4vLi4vLi4vLi4vLi4vLi4vdGZqcy1iYWNrZW5kLXdhc20vc3JjL2tlcm5lbHMvRGVwdGhUb1NwYWNlLnRzIl0sIm5hbWVzIjpbXSwibWFwcGluZ3MiOiJBQUFBOzs7Ozs7Ozs7Ozs7Ozs7R0FlRztBQUVILE9BQU8sRUFBQyxZQUFZLEVBQStFLElBQUksRUFBQyxNQUFNLHVCQUF1QixDQUFDO0FBSXRJLElBQUksZ0JBR3VDLENBQUM7QUFFNUMsU0FBUyxLQUFLLENBQUMsT0FBb0I7SUFDakMsZ0JBQWdCLEdBQUcsT0FBTyxDQUFDLElBQUksQ0FBQyxLQUFLLENBQUMsWUFBWSxFQUFFLElBQUksQ0FBQyxRQUFRLEVBQUU7UUFDakUsUUFBUTtRQUNSLFFBQVE7UUFDUixRQUFRO1FBQ1IsT0FBTztRQUNQLFFBQVE7UUFDUixPQUFPO1FBQ1AsT0FBTztRQUNQLFFBQVE7UUFDUixRQUFRLEVBQUcsUUFBUTtLQUNwQixDQUFDLENBQUM7QUFDTCxDQUFDO0FBRUQsTUFBTSxVQUFVLFlBQVksQ0FBQyxJQUk1QjtJQUNDLE1BQU0sRUFBQyxPQUFPLEVBQUUsTUFBTSxFQUFFLEtBQUssRUFBQyxHQUFHLElBQUksQ0FBQztJQUN0QyxNQUFNLEVBQUMsQ0FBQyxFQUFDLEdBQUcsTUFBTSxDQUFDO0lBQ25CLE1BQU0sRUFBQyxTQUFTLEVBQUUsVUFBVSxFQUFDLEdBQUcsS0FBSyxDQUFDO0lBRXRDLE1BQU0sU0FBUyxHQUFHLENBQUMsQ0FBQyxLQUFLLENBQUMsQ0FBQyxDQUFDLENBQUM7SUFDN0IsTUFBTSxXQUFXLEdBQUcsQ0FBQyxVQUFVLEtBQUssTUFBTSxDQUFDLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxLQUFLLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxLQUFLLENBQUMsQ0FBQyxDQUFDLENBQUM7SUFDdEUsTUFBTSxVQUFVLEdBQUcsQ0FBQyxVQUFVLEtBQUssTUFBTSxDQUFDLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxLQUFLLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxLQUFLLENBQUMsQ0FBQyxDQUFDLENBQUM7SUFDckUsTUFBTSxVQUFVLEdBQUcsQ0FBQyxVQUFVLEtBQUssTUFBTSxDQUFDLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxLQUFLLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxLQUFLLENBQUMsQ0FBQyxDQUFDLENBQUM7SUFFckUsTUFBTSxZQUFZLEdBQUcsV0FBVyxHQUFHLFNBQVMsQ0FBQztJQUM3QyxNQUFNLFdBQVcsR0FBRyxVQUFVLEdBQUcsU0FBUyxDQUFDO0lBQzNDLE1BQU0sV0FBVyxHQUFHLFVBQVUsR0FBRyxDQUFDLFNBQVMsR0FBRyxTQUFTLENBQUMsQ0FBQztJQUV6RCxNQUFNLFdBQVcsR0FBRyxDQUFDLFVBQVUsS0FBSyxNQUFNLENBQUMsQ0FBQyxDQUFDO1FBQ3pDLENBQUMsU0FBUyxFQUFFLFlBQVksRUFBRSxXQUFXLEVBQUUsV0FBVyxDQUFDLENBQUMsQ0FBQztRQUNyRCxDQUFDLFNBQVMsRUFBRSxXQUFXLEVBQUUsWUFBWSxFQUFFLFdBQVcsQ0FBQyxDQUFDO0lBRXhELE1BQU0sR0FBRyxHQUFHLE9BQU8sQ0FBQyxVQUFVLENBQUMsV0FBVyxFQUFFLFNBQVMsQ0FBQyxDQUFDO0lBRXZELE1BQU0sS0FBSyxHQUFHLE9BQU8sQ0FBQyxTQUFTLENBQUMsR0FBRyxDQUFDLENBQUMsQ0FBQyxNQUFNLENBQUMsQ0FBQztJQUM5QyxNQUFNLEdBQUcsR0FBRyxLQUFLLENBQUMsRUFBRSxDQUFDO0lBQ3JCLE1BQU0sYUFBYSxHQUNmLElBQUksVUFBVSxDQUFDLElBQUksVUFBVSxDQUFDLElBQUksQ0FBQyxjQUFjLENBQUMsQ0FBQyxDQUFDLEtBQUssQ0FBQyxDQUFDLENBQUMsTUFBTSxDQUFDLENBQUM7SUFFeEUsTUFBTSxnQkFBZ0IsR0FBRyxJQUFJLFVBQVUsQ0FBQyxJQUFJLFVBQVUsQ0FBQyxXQUFXLENBQUMsQ0FBQyxNQUFNLENBQUMsQ0FBQztJQUM1RSxNQUFNLGVBQWUsR0FDakIsSUFBSSxVQUFVLENBQUMsSUFBSSxVQUFVLENBQUMsSUFBSSxDQUFDLGNBQWMsQ0FBQyxXQUFXLENBQUMsQ0FBQyxDQUFDLE1BQU0sQ0FBQyxDQUFDO0lBRTVFLE1BQU0sS0FBSyxHQUFHLE9BQU8sQ0FBQyxTQUFTLENBQUMsR0FBRyxDQUFDLEdBQUcsQ0FBQyxNQUFNLENBQUMsQ0FBQyxFQUFFLENBQUM7SUFDbkQsTUFBTSxZQUFZLEdBQUcsVUFBVSxLQUFLLE1BQU0sQ0FBQyxDQUFDLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxDQUFDLENBQUM7SUFDbkQsZ0JBQWdCLENBQ1osR0FBRyxFQUFFLFNBQVMsRUFBRSxZQUFZLEVBQUUsYUFBYSxFQUFFLENBQUMsQ0FBQyxLQUFLLENBQUMsTUFBTSxHQUFHLENBQUMsRUFDL0QsZ0JBQWdCLEVBQUUsZUFBZSxFQUFFLFdBQVcsQ0FBQyxNQUFNLEVBQUUsS0FBSyxDQUFDLENBQUM7SUFFbEUsT0FBTyxHQUFHLENBQUM7QUFDYixDQUFDO0FBRUQsTUFBTSxDQUFDLE1BQU0sa0JBQWtCLEdBQWlCO0lBQzlDLFVBQVUsRUFBRSxZQUFZO0lBQ3hCLFdBQVcsRUFBRSxNQUFNO0lBQ25CLFNBQVMsRUFBRSxLQUFLO0lBQ2hCLFVBQVUsRUFBRSxZQUFxQztDQUNsRCxDQUFDIiwic291cmNlc0NvbnRlbnQiOlsiLyoqXG4gKiBAbGljZW5zZVxuICogQ29weXJpZ2h0IDIwMTkgR29vZ2xlIExMQy4gQWxsIFJpZ2h0cyBSZXNlcnZlZC5cbiAqIExpY2Vuc2VkIHVuZGVyIHRoZSBBcGFjaGUgTGljZW5zZSwgVmVyc2lvbiAyLjAgKHRoZSBcIkxpY2Vuc2VcIik7XG4gKiB5b3UgbWF5IG5vdCB1c2UgdGhpcyBmaWxlIGV4Y2VwdCBpbiBjb21wbGlhbmNlIHdpdGggdGhlIExpY2Vuc2UuXG4gKiBZb3UgbWF5IG9idGFpbiBhIGNvcHkgb2YgdGhlIExpY2Vuc2UgYXRcbiAqXG4gKiBodHRwOi8vd3d3LmFwYWNoZS5vcmcvbGljZW5zZXMvTElDRU5TRS0yLjBcbiAqXG4gKiBVbmxlc3MgcmVxdWlyZWQgYnkgYXBwbGljYWJsZSBsYXcgb3IgYWdyZWVkIHRvIGluIHdyaXRpbmcsIHNvZnR3YXJlXG4gKiBkaXN0cmlidXRlZCB1bmRlciB0aGUgTGljZW5zZSBpcyBkaXN0cmlidXRlZCBvbiBhbiBcIkFTIElTXCIgQkFTSVMsXG4gKiBXSVRIT1VUIFdBUlJBTlRJRVMgT1IgQ09ORElUSU9OUyBPRiBBTlkgS0lORCwgZWl0aGVyIGV4cHJlc3Mgb3IgaW1wbGllZC5cbiAqIFNlZSB0aGUgTGljZW5zZSBmb3IgdGhlIHNwZWNpZmljIGxhbmd1YWdlIGdvdmVybmluZyBwZXJtaXNzaW9ucyBhbmRcbiAqIGxpbWl0YXRpb25zIHVuZGVyIHRoZSBMaWNlbnNlLlxuICogPT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT1cbiAqL1xuXG5pbXBvcnQge0RlcHRoVG9TcGFjZSwgRGVwdGhUb1NwYWNlQXR0cnMsIERlcHRoVG9TcGFjZUlucHV0cywgS2VybmVsQ29uZmlnLCBLZXJuZWxGdW5jLCBUZW5zb3JJbmZvLCB1dGlsfSBmcm9tICdAdGVuc29yZmxvdy90ZmpzLWNvcmUnO1xuXG5pbXBvcnQge0JhY2tlbmRXYXNtfSBmcm9tICcuLi9iYWNrZW5kX3dhc20nO1xuXG5sZXQgd2FzbURlcHRoVG9TcGFjZTogKFxuICAgIHhJZDogbnVtYmVyLCBibG9ja1NpemU6IG51bWJlciwgY2hhbm5lbHNMYXN0OiBudW1iZXIsIHhTdHJpZGVzOiBVaW50OEFycmF5LFxuICAgIHhTdHJpZGVzTGVuZ3RoOiBudW1iZXIsIG91dHB1dFNoYXBlOiBVaW50OEFycmF5LCBvdXRwdXRTdHJpZGVzOiBVaW50OEFycmF5LFxuICAgIG91dFNpemU6IG51bWJlciwgb3V0SWQ6IG51bWJlcikgPT4gdm9pZDtcblxuZnVuY3Rpb24gc2V0dXAoYmFja2VuZDogQmFja2VuZFdhc20pOiB2b2lkIHtcbiAgd2FzbURlcHRoVG9TcGFjZSA9IGJhY2tlbmQud2FzbS5jd3JhcChEZXB0aFRvU3BhY2UsIG51bGwgLyp2b2lkKi8sIFtcbiAgICAnbnVtYmVyJywgIC8vIHhJZFxuICAgICdudW1iZXInLCAgLy8gYmxvY2tTaXplXG4gICAgJ251bWJlcicsICAvLyBjaGFubmVsc0xhc3RcbiAgICAnYXJyYXknLCAgIC8vIHhTdHJpZGVzXG4gICAgJ251bWJlcicsICAvLyB4U3RyaWRlc0xlbmd0aFxuICAgICdhcnJheScsICAgLy8gb3V0cHV0U2hhcGVcbiAgICAnYXJyYXknLCAgIC8vIG91dHB1dFN0cmlkZXNcbiAgICAnbnVtYmVyJywgIC8vIG91dFNpemVcbiAgICAnbnVtYmVyJywgIC8vIG91dElkXG4gIF0pO1xufVxuXG5leHBvcnQgZnVuY3Rpb24gZGVwdGhUb1NwYWNlKGFyZ3M6IHtcbiAgYmFja2VuZDogQmFja2VuZFdhc20sXG4gIGlucHV0czogRGVwdGhUb1NwYWNlSW5wdXRzLFxuICBhdHRyczogRGVwdGhUb1NwYWNlQXR0cnNcbn0pOiBUZW5zb3JJbmZvIHtcbiAgY29uc3Qge2JhY2tlbmQsIGlucHV0cywgYXR0cnN9ID0gYXJncztcbiAgY29uc3Qge3h9ID0gaW5wdXRzO1xuICBjb25zdCB7YmxvY2tTaXplLCBkYXRhRm9ybWF0fSA9IGF0dHJzO1xuXG4gIGNvbnN0IGJhdGNoU2l6ZSA9IHguc2hhcGVbMF07XG4gIGNvbnN0IGlucHV0SGVpZ2h0ID0gKGRhdGFGb3JtYXQgPT09ICdOSFdDJykgPyB4LnNoYXBlWzFdIDogeC5zaGFwZVsyXTtcbiAgY29uc3QgaW5wdXRXaWR0aCA9IChkYXRhRm9ybWF0ID09PSAnTkhXQycpID8geC5zaGFwZVsyXSA6IHguc2hhcGVbM107XG4gIGNvbnN0IGlucHV0RGVwdGggPSAoZGF0YUZvcm1hdCA9PT0gJ05IV0MnKSA/IHguc2hhcGVbM10gOiB4LnNoYXBlWzFdO1xuXG4gIGNvbnN0IG91dHB1dEhlaWdodCA9IGlucHV0SGVpZ2h0ICogYmxvY2tTaXplO1xuICBjb25zdCBvdXRwdXRXaWR0aCA9IGlucHV0V2lkdGggKiBibG9ja1NpemU7XG4gIGNvbnN0IG91dHB1dERlcHRoID0gaW5wdXREZXB0aCAvIChibG9ja1NpemUgKiBibG9ja1NpemUpO1xuXG4gIGNvbnN0IG91dHB1dFNoYXBlID0gKGRhdGFGb3JtYXQgPT09ICdOSFdDJykgP1xuICAgICAgW2JhdGNoU2l6ZSwgb3V0cHV0SGVpZ2h0LCBvdXRwdXRXaWR0aCwgb3V0cHV0RGVwdGhdIDpcbiAgICAgIFtiYXRjaFNpemUsIG91dHB1dERlcHRoLCBvdXRwdXRIZWlnaHQsIG91dHB1dFdpZHRoXTtcblxuICBjb25zdCBvdXQgPSBiYWNrZW5kLm1ha2VPdXRwdXQob3V0cHV0U2hhcGUsICdmbG9hdDMyJyk7XG5cbiAgY29uc3QgeERhdGEgPSBiYWNrZW5kLmRhdGFJZE1hcC5nZXQoeC5kYXRhSWQpO1xuICBjb25zdCB4SWQgPSB4RGF0YS5pZDtcbiAgY29uc3QgeFN0cmlkZXNCeXRlcyA9XG4gICAgICBuZXcgVWludDhBcnJheShuZXcgSW50MzJBcnJheSh1dGlsLmNvbXB1dGVTdHJpZGVzKHguc2hhcGUpKS5idWZmZXIpO1xuXG4gIGNvbnN0IG91dHB1dFNoYXBlQnl0ZXMgPSBuZXcgVWludDhBcnJheShuZXcgSW50MzJBcnJheShvdXRwdXRTaGFwZSkuYnVmZmVyKTtcbiAgY29uc3Qgb3V0U3RyaWRlc0J5dGVzID1cbiAgICAgIG5ldyBVaW50OEFycmF5KG5ldyBJbnQzMkFycmF5KHV0aWwuY29tcHV0ZVN0cmlkZXMob3V0cHV0U2hhcGUpKS5idWZmZXIpO1xuXG4gIGNvbnN0IG91dElkID0gYmFja2VuZC5kYXRhSWRNYXAuZ2V0KG91dC5kYXRhSWQpLmlkO1xuICBjb25zdCBjaGFubmVsc0xhc3QgPSBkYXRhRm9ybWF0ID09PSAnTkhXQycgPyAxIDogMDtcbiAgd2FzbURlcHRoVG9TcGFjZShcbiAgICAgIHhJZCwgYmxvY2tTaXplLCBjaGFubmVsc0xhc3QsIHhTdHJpZGVzQnl0ZXMsIHguc2hhcGUubGVuZ3RoIC0gMSxcbiAgICAgIG91dHB1dFNoYXBlQnl0ZXMsIG91dFN0cmlkZXNCeXRlcywgb3V0cHV0U2hhcGUubGVuZ3RoLCBvdXRJZCk7XG5cbiAgcmV0dXJuIG91dDtcbn1cblxuZXhwb3J0IGNvbnN0IGRlcHRoVG9TcGFjZUNvbmZpZzogS2VybmVsQ29uZmlnID0ge1xuICBrZXJuZWxOYW1lOiBEZXB0aFRvU3BhY2UsXG4gIGJhY2tlbmROYW1lOiAnd2FzbScsXG4gIHNldHVwRnVuYzogc2V0dXAsXG4gIGtlcm5lbEZ1bmM6IGRlcHRoVG9TcGFjZSBhcyB1bmtub3duIGFzIEtlcm5lbEZ1bmNcbn07XG4iXX0=