From 8cfdc86256139d68b21337a227dafa0898546893 Mon Sep 17 00:00:00 2001 From: Ann Yuan Date: Tue, 18 Feb 2020 16:03:19 -0500 Subject: [PATCH 1/4] add kernel --- tfjs-node/src/nodejs_kernel_backend.ts | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index 0545aee0fdc..9afc0309886 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -449,6 +449,19 @@ export class NodeJSKernelBackend extends KernelBackend { return this.executeSingleOutput('AddN', opAttrs, tensors) as T; } + softmax(logits: T, dim: number): T { + const axes = util.parseAxisParam([dim], logits.shape); + + const maxLogit = this.max(logits, axes); + const expandedShape = + backend_util.expandShapeToKeepDim(maxLogit.shape, axes); + const a = this.subtract(logits, maxLogit.reshape(expandedShape)); + const b = this.exp(a); + const sumExp = this.sum(b, axes).reshape(expandedShape); + + return this.realDivide(b, sumExp) as T; + } + subtract(a: Tensor, b: Tensor): Tensor { const opAttrs = [createTypeOpAttr('T', backend_util.upcastType(a.dtype, b.dtype))]; From 1683bf77bd8d240d25c0c931a8f2a8a6d32ba650 Mon Sep 17 00:00:00 2001 From: Ann Yuan Date: Tue, 18 Feb 2020 16:41:21 -0500 Subject: [PATCH 2/4] add --- tfjs-node/src/nodejs_kernel_backend.ts | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index 9afc0309886..34a7a3e2c6b 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -450,16 +450,8 @@ export class NodeJSKernelBackend extends KernelBackend { } softmax(logits: T, dim: number): T { - const axes = util.parseAxisParam([dim], logits.shape); - - const maxLogit = this.max(logits, axes); - const expandedShape = - backend_util.expandShapeToKeepDim(maxLogit.shape, axes); - const a = this.subtract(logits, maxLogit.reshape(expandedShape)); - const b = this.exp(a); - const sumExp = this.sum(b, axes).reshape(expandedShape); - - return this.realDivide(b, sumExp) as T; + const opAttrs = [createTypeOpAttr('T', logits.dtype)]; + return this.executeSingleOutput('Softmax', opAttrs, [logits]) as T; } subtract(a: Tensor, b: Tensor): Tensor { From f6bb65e42ff297e1dc4724022d8faf747b4a2a55 Mon Sep 17 00:00:00 2001 From: Ann Yuan Date: Tue, 18 Feb 2020 16:45:06 -0500 Subject: [PATCH 3/4] add modular kernel --- tfjs-node/src/kernels/Softmax.ts | 38 ++++++++++++++++++++++++++ tfjs-node/src/kernels/all_kernels.ts | 1 + tfjs-node/src/nodejs_kernel_backend.ts | 5 ---- 3 files changed, 39 insertions(+), 5 deletions(-) create mode 100644 tfjs-node/src/kernels/Softmax.ts diff --git a/tfjs-node/src/kernels/Softmax.ts b/tfjs-node/src/kernels/Softmax.ts new file mode 100644 index 00000000000..215125ff91d --- /dev/null +++ b/tfjs-node/src/kernels/Softmax.ts @@ -0,0 +1,38 @@ +/** + * @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 {NamedTensorInfoMap, registerKernel, TensorInfo} from '@tensorflow/tfjs-core'; + +import {createTypeOpAttr, NodeJSKernelBackend} from '../nodejs_kernel_backend'; + +interface SoftmaxInputs extends NamedTensorInfoMap { + logits: TensorInfo; +} + +// TODO(nsthorat, dsmilkov): Remove dependency on tensors, use dataId. +registerKernel({ + kernelName: 'Softmax', + backendName: 'tensorflow', + kernelFunc: ({inputs, backend}) => { + const {logits} = inputs as SoftmaxInputs; + const opAttrs = [createTypeOpAttr('T', logits.dtype)]; + + const nodeBackend = backend as NodeJSKernelBackend; + + return nodeBackend.executeSingleOutput('Softmax', opAttrs, [logits]); + } +}); diff --git a/tfjs-node/src/kernels/all_kernels.ts b/tfjs-node/src/kernels/all_kernels.ts index c6eea842964..6b5f6bf8245 100644 --- a/tfjs-node/src/kernels/all_kernels.ts +++ b/tfjs-node/src/kernels/all_kernels.ts @@ -19,3 +19,4 @@ // 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'; +import './Softmax'; diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index 34a7a3e2c6b..0545aee0fdc 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -449,11 +449,6 @@ export class NodeJSKernelBackend extends KernelBackend { return this.executeSingleOutput('AddN', opAttrs, tensors) as T; } - softmax(logits: T, dim: number): T { - const opAttrs = [createTypeOpAttr('T', logits.dtype)]; - return this.executeSingleOutput('Softmax', opAttrs, [logits]) as T; - } - subtract(a: Tensor, b: Tensor): Tensor { const opAttrs = [createTypeOpAttr('T', backend_util.upcastType(a.dtype, b.dtype))]; From a63f5a6bdf4e3b40ed33e36ceaaf249fc3632619 Mon Sep 17 00:00:00 2001 From: Ann Yuan Date: Tue, 18 Feb 2020 16:46:50 -0500 Subject: [PATCH 4/4] remove comment --- tfjs-node/src/kernels/Softmax.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/tfjs-node/src/kernels/Softmax.ts b/tfjs-node/src/kernels/Softmax.ts index 215125ff91d..af08a9f15d0 100644 --- a/tfjs-node/src/kernels/Softmax.ts +++ b/tfjs-node/src/kernels/Softmax.ts @@ -23,7 +23,6 @@ interface SoftmaxInputs extends NamedTensorInfoMap { logits: TensorInfo; } -// TODO(nsthorat, dsmilkov): Remove dependency on tensors, use dataId. registerKernel({ kernelName: 'Softmax', backendName: 'tensorflow',