From ffd7172694fcae1cfc099c91e7ab4a626617e801 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 19 Dec 2019 13:37:56 -0800 Subject: [PATCH 1/6] Add support for NonMaxSuppressionV5 using modular pattern for Node --- tfjs-node/src/all_kernels.ts | 21 ++++++++++ tfjs-node/src/index.ts | 3 ++ tfjs-node/src/nodejs_kernel_backend.ts | 4 +- tfjs-node/src/non_max_suppression_v5.ts | 52 +++++++++++++++++++++++++ 4 files changed, 78 insertions(+), 2 deletions(-) create mode 100644 tfjs-node/src/all_kernels.ts create mode 100644 tfjs-node/src/non_max_suppression_v5.ts diff --git a/tfjs-node/src/all_kernels.ts b/tfjs-node/src/all_kernels.ts new file mode 100644 index 00000000000..c6eea842964 --- /dev/null +++ b/tfjs-node/src/all_kernels.ts @@ -0,0 +1,21 @@ +/** + * @license + * Copyright 2019 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. + * ============================================================================= + */ + +// We explicitly import the modular kernels so they get registered in the +// global registry when we compile the library. A modular build would replace +// the contents of this file and import only the kernels that are needed. +import './non_max_suppression_v5'; diff --git a/tfjs-node/src/index.ts b/tfjs-node/src/index.ts index 2ca6bbf87fd..3dd11de284f 100644 --- a/tfjs-node/src/index.ts +++ b/tfjs-node/src/index.ts @@ -24,6 +24,9 @@ import {NodeJSKernelBackend} from './nodejs_kernel_backend'; import {TFJSBinding} from './tfjs_binding'; import * as nodeVersion from './version'; +// Import all kernels. +import './all_kernels'; + // tslint:disable-next-line:no-require-imports const binary = require('node-pre-gyp'); const bindingPath = diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index cca6fd2e229..23a30a26d56 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -68,7 +68,7 @@ export class NodeJSKernelBackend extends KernelBackend { } // Creates a new Tensor and maps the dataId to the passed in ID. - private createOutputTensor(metadata: TensorMetadata): Tensor { + createOutputTensor(metadata: TensorMetadata): Tensor { const newId = {}; this.tensorMap.set(newId, { @@ -112,7 +112,7 @@ export class NodeJSKernelBackend extends KernelBackend { } // Prepares Tensor instances for Op execution. - private getInputTensorIds(tensors: Array): number[] { + getInputTensorIds(tensors: Array): number[] { const ids: number[] = []; for (let i = 0; i < tensors.length; i++) { if (tensors[i] instanceof Int64Scalar) { diff --git a/tfjs-node/src/non_max_suppression_v5.ts b/tfjs-node/src/non_max_suppression_v5.ts new file mode 100644 index 00000000000..24a4a2925c6 --- /dev/null +++ b/tfjs-node/src/non_max_suppression_v5.ts @@ -0,0 +1,52 @@ +/** + * @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 {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, scalar} from '@tensorflow/tfjs-core'; +import { NodeJSKernelBackend, createTypeOpAttr } from './nodejs_kernel_backend'; + +interface NonMaxSuppressionWithScoreInputs extends NamedTensorInfoMap { + boxes: TensorInfo; + scores: TensorInfo; +} + +interface NonMaxSuppressionWithScoreAttrs extends NamedAttrMap { + maxOutputSize: number; + iouThreshold: number; + scoreThreshold: number; + softNmsSigma: number; +} + +registerKernel({ + kernelName: 'NonMaxSuppressionV5', + backendName: 'tensorflow', + kernelFunc: ({inputs, backend, attrs}) => { + const {boxes, scores} = inputs as NonMaxSuppressionWithScoreInputs; + const {maxOutputSize, iouThreshold, scoreThreshold, softNmsSigma} = + attrs as NonMaxSuppressionWithScoreAttrs; + const maxOutputSizeTensor = scalar(maxOutputSize, 'int32'); + const iouThresholdTensor = scalar(iouThreshold); + const scoreThresholdTensor = scalar(scoreThreshold); + const softNmsSigmaTensor = scalar(softNmsSigma); + const opAttrs = [createTypeOpAttr('T', boxes.dtype)]; + + const nodeBackend = backend as NodeJSKernelBackend; + + const outputMetadata = nodeBackend.binding.executeOp('NonMaxSuppressionV5', opAttrs, nodeBackend.getInputTensorIds([boxes, scores, maxOutputSizeTensor, iouThresholdTensor, scoreThresholdTensor, softNmsSigmaTensor]), 2); + + return outputMetadata.map(m => nodeBackend.createOutputTensor(m)); + } +}); From 57134a2979a9ac06181af44fc0b0d0eb824588e3 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 19 Dec 2019 13:46:53 -0800 Subject: [PATCH 2/6] Create a kernels directory for all new kernels --- tfjs-node/src/index.ts | 2 +- tfjs-node/src/{ => kernels}/all_kernels.ts | 0 tfjs-node/src/{ => kernels}/non_max_suppression_v5.ts | 2 +- tfjs-node/src/nodejs_kernel_backend.ts | 6 +++--- 4 files changed, 5 insertions(+), 5 deletions(-) rename tfjs-node/src/{ => kernels}/all_kernels.ts (100%) rename tfjs-node/src/{ => kernels}/non_max_suppression_v5.ts (96%) diff --git a/tfjs-node/src/index.ts b/tfjs-node/src/index.ts index 3dd11de284f..6533cde569f 100644 --- a/tfjs-node/src/index.ts +++ b/tfjs-node/src/index.ts @@ -25,7 +25,7 @@ import {TFJSBinding} from './tfjs_binding'; import * as nodeVersion from './version'; // Import all kernels. -import './all_kernels'; +import './kernels/all_kernels'; // tslint:disable-next-line:no-require-imports const binary = require('node-pre-gyp'); diff --git a/tfjs-node/src/all_kernels.ts b/tfjs-node/src/kernels/all_kernels.ts similarity index 100% rename from tfjs-node/src/all_kernels.ts rename to tfjs-node/src/kernels/all_kernels.ts diff --git a/tfjs-node/src/non_max_suppression_v5.ts b/tfjs-node/src/kernels/non_max_suppression_v5.ts similarity index 96% rename from tfjs-node/src/non_max_suppression_v5.ts rename to tfjs-node/src/kernels/non_max_suppression_v5.ts index 24a4a2925c6..f58640bd867 100644 --- a/tfjs-node/src/non_max_suppression_v5.ts +++ b/tfjs-node/src/kernels/non_max_suppression_v5.ts @@ -16,7 +16,7 @@ */ import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, scalar} from '@tensorflow/tfjs-core'; -import { NodeJSKernelBackend, createTypeOpAttr } from './nodejs_kernel_backend'; +import { NodeJSKernelBackend, createTypeOpAttr } from '../nodejs_kernel_backend'; interface NonMaxSuppressionWithScoreInputs extends NamedTensorInfoMap { boxes: TensorInfo; diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index 23a30a26d56..294771aa989 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -46,7 +46,7 @@ export class NodeJSKernelBackend extends KernelBackend { this.tensorMap = new tfc.DataStorage(this, tfc.engine()); } - private getDTypeInteger(dtype: DataType): number { + getDTypeInteger(dtype: DataType): number { switch (dtype) { case 'float32': return this.binding.TF_FLOAT; @@ -63,7 +63,7 @@ export class NodeJSKernelBackend extends KernelBackend { } } - private typeAttributeFromTensor(value: Tensor): number { + typeAttributeFromTensor(value: Tensor): number { return this.getDTypeInteger(value.dtype); } @@ -137,7 +137,7 @@ export class NodeJSKernelBackend extends KernelBackend { return ids; } - private createReductionOpAttrs(tensor: Tensor): TFEOpAttr[] { + createReductionOpAttrs(tensor: Tensor): TFEOpAttr[] { return [ {name: 'keep_dims', type: this.binding.TF_ATTR_BOOL, value: false}, createTypeOpAttr('T', tensor.dtype), createTypeOpAttr('Tidx', 'int32') From 528faccee666fbe144a0c981588088569bc79ae1 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 19 Dec 2019 14:04:15 -0800 Subject: [PATCH 3/6] Fix lint --- .../src/kernels/non_max_suppression_v5.ts | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tfjs-node/src/kernels/non_max_suppression_v5.ts b/tfjs-node/src/kernels/non_max_suppression_v5.ts index f58640bd867..ddf828c073d 100644 --- a/tfjs-node/src/kernels/non_max_suppression_v5.ts +++ b/tfjs-node/src/kernels/non_max_suppression_v5.ts @@ -15,8 +15,9 @@ * ============================================================================= */ -import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, scalar} from '@tensorflow/tfjs-core'; -import { NodeJSKernelBackend, createTypeOpAttr } from '../nodejs_kernel_backend'; +import {NamedAttrMap, NamedTensorInfoMap, registerKernel, scalar, TensorInfo} from '@tensorflow/tfjs-core'; + +import {createTypeOpAttr, NodeJSKernelBackend} from '../nodejs_kernel_backend'; interface NonMaxSuppressionWithScoreInputs extends NamedTensorInfoMap { boxes: TensorInfo; @@ -37,15 +38,20 @@ registerKernel({ const {boxes, scores} = inputs as NonMaxSuppressionWithScoreInputs; const {maxOutputSize, iouThreshold, scoreThreshold, softNmsSigma} = attrs as NonMaxSuppressionWithScoreAttrs; - const maxOutputSizeTensor = scalar(maxOutputSize, 'int32'); - const iouThresholdTensor = scalar(iouThreshold); - const scoreThresholdTensor = scalar(scoreThreshold); - const softNmsSigmaTensor = scalar(softNmsSigma); + const maxOutputSizeTensor = scalar(maxOutputSize, 'int32'); + const iouThresholdTensor = scalar(iouThreshold); + const scoreThresholdTensor = scalar(scoreThreshold); + const softNmsSigmaTensor = scalar(softNmsSigma); const opAttrs = [createTypeOpAttr('T', boxes.dtype)]; const nodeBackend = backend as NodeJSKernelBackend; - const outputMetadata = nodeBackend.binding.executeOp('NonMaxSuppressionV5', opAttrs, nodeBackend.getInputTensorIds([boxes, scores, maxOutputSizeTensor, iouThresholdTensor, scoreThresholdTensor, softNmsSigmaTensor]), 2); + const outputMetadata = nodeBackend.binding.executeOp( + 'NonMaxSuppressionV5', opAttrs, nodeBackend.getInputTensorIds([ + boxes, scores, maxOutputSizeTensor, iouThresholdTensor, + scoreThresholdTensor, softNmsSigmaTensor + ]), + 2); return outputMetadata.map(m => nodeBackend.createOutputTensor(m)); } From 10f49b896b285683fa72ae2e7e43a9cbbdffbfd6 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 19 Dec 2019 14:16:14 -0800 Subject: [PATCH 4/6] Use public method instead --- tfjs-node/src/kernels/non_max_suppression_v5.ts | 15 +++++++-------- tfjs-node/src/nodejs_kernel_backend.ts | 10 +++++----- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/tfjs-node/src/kernels/non_max_suppression_v5.ts b/tfjs-node/src/kernels/non_max_suppression_v5.ts index ddf828c073d..8cf59c8d08d 100644 --- a/tfjs-node/src/kernels/non_max_suppression_v5.ts +++ b/tfjs-node/src/kernels/non_max_suppression_v5.ts @@ -15,7 +15,7 @@ * ============================================================================= */ -import {NamedAttrMap, NamedTensorInfoMap, registerKernel, scalar, TensorInfo} from '@tensorflow/tfjs-core'; +import {NamedAttrMap, NamedTensorInfoMap, registerKernel, scalar, Tensor1D, Tensor2D, TensorInfo} from '@tensorflow/tfjs-core'; import {createTypeOpAttr, NodeJSKernelBackend} from '../nodejs_kernel_backend'; @@ -46,13 +46,12 @@ registerKernel({ const nodeBackend = backend as NodeJSKernelBackend; - const outputMetadata = nodeBackend.binding.executeOp( - 'NonMaxSuppressionV5', opAttrs, nodeBackend.getInputTensorIds([ - boxes, scores, maxOutputSizeTensor, iouThresholdTensor, - scoreThresholdTensor, softNmsSigmaTensor - ]), + return nodeBackend.executeMultipleOutputs( + 'NonMaxSuppressionV5', opAttrs, + [ + boxes as Tensor2D, scores as Tensor1D, maxOutputSizeTensor, + iouThresholdTensor, scoreThresholdTensor, softNmsSigmaTensor + ], 2); - - return outputMetadata.map(m => nodeBackend.createOutputTensor(m)); } }); diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index 294771aa989..cca6fd2e229 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -46,7 +46,7 @@ export class NodeJSKernelBackend extends KernelBackend { this.tensorMap = new tfc.DataStorage(this, tfc.engine()); } - getDTypeInteger(dtype: DataType): number { + private getDTypeInteger(dtype: DataType): number { switch (dtype) { case 'float32': return this.binding.TF_FLOAT; @@ -63,12 +63,12 @@ export class NodeJSKernelBackend extends KernelBackend { } } - typeAttributeFromTensor(value: Tensor): number { + private typeAttributeFromTensor(value: Tensor): number { return this.getDTypeInteger(value.dtype); } // Creates a new Tensor and maps the dataId to the passed in ID. - createOutputTensor(metadata: TensorMetadata): Tensor { + private createOutputTensor(metadata: TensorMetadata): Tensor { const newId = {}; this.tensorMap.set(newId, { @@ -112,7 +112,7 @@ export class NodeJSKernelBackend extends KernelBackend { } // Prepares Tensor instances for Op execution. - getInputTensorIds(tensors: Array): number[] { + private getInputTensorIds(tensors: Array): number[] { const ids: number[] = []; for (let i = 0; i < tensors.length; i++) { if (tensors[i] instanceof Int64Scalar) { @@ -137,7 +137,7 @@ export class NodeJSKernelBackend extends KernelBackend { return ids; } - createReductionOpAttrs(tensor: Tensor): TFEOpAttr[] { + private createReductionOpAttrs(tensor: Tensor): TFEOpAttr[] { return [ {name: 'keep_dims', type: this.binding.TF_ATTR_BOOL, value: false}, createTypeOpAttr('T', tensor.dtype), createTypeOpAttr('Tidx', 'int32') From 3a85ea891bec6a7672623aabb1582608d782aab2 Mon Sep 17 00:00:00 2001 From: Na Li Date: Fri, 20 Dec 2019 11:20:42 -0800 Subject: [PATCH 5/6] Solve memory leak --- .../src/kernels/non_max_suppression_v5.ts | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/tfjs-node/src/kernels/non_max_suppression_v5.ts b/tfjs-node/src/kernels/non_max_suppression_v5.ts index 8cf59c8d08d..a1c1cda27fa 100644 --- a/tfjs-node/src/kernels/non_max_suppression_v5.ts +++ b/tfjs-node/src/kernels/non_max_suppression_v5.ts @@ -46,12 +46,21 @@ registerKernel({ const nodeBackend = backend as NodeJSKernelBackend; - return nodeBackend.executeMultipleOutputs( - 'NonMaxSuppressionV5', opAttrs, - [ - boxes as Tensor2D, scores as Tensor1D, maxOutputSizeTensor, - iouThresholdTensor, scoreThresholdTensor, softNmsSigmaTensor - ], - 2); + const [selectedIndices, selectedScores, validOutputs] = + nodeBackend.executeMultipleOutputs( + 'NonMaxSuppressionV5', opAttrs, + [ + boxes as Tensor2D, scores as Tensor1D, maxOutputSizeTensor, + iouThresholdTensor, scoreThresholdTensor, softNmsSigmaTensor + ], + 3); + + maxOutputSizeTensor.dispose(); + iouThresholdTensor.dispose(); + scoreThresholdTensor.dispose(); + softNmsSigmaTensor.dispose(); + validOutputs.dispose(); + + return [selectedIndices, selectedScores]; } }); From 312bc1582143476067dd65a1f4b8f298dd5b9273 Mon Sep 17 00:00:00 2001 From: Na Li Date: Fri, 20 Dec 2019 14:10:50 -0800 Subject: [PATCH 6/6] Add TODO --- tfjs-node/src/kernels/non_max_suppression_v5.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/tfjs-node/src/kernels/non_max_suppression_v5.ts b/tfjs-node/src/kernels/non_max_suppression_v5.ts index a1c1cda27fa..6f6dbb46b89 100644 --- a/tfjs-node/src/kernels/non_max_suppression_v5.ts +++ b/tfjs-node/src/kernels/non_max_suppression_v5.ts @@ -31,6 +31,7 @@ interface NonMaxSuppressionWithScoreAttrs extends NamedAttrMap { softNmsSigma: number; } +// TODO(nsthorat, dsmilkov): Remove dependency on tensors, use dataId. registerKernel({ kernelName: 'NonMaxSuppressionV5', backendName: 'tensorflow',