From 6c9bd4512936671606462c24f2a5bbb6c2209390 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 24 Aug 2020 13:00:29 -0700 Subject: [PATCH 01/22] Resolve merge conflict. --- tfjs-backend-cpu/src/backend_cpu.ts | 52 +++++++++++++++----- tfjs-backend-cpu/src/kernels/Reshape.ts | 34 +++++++++++++ tfjs-backend-cpu/src/kernels/Reshape_impl.ts | 31 ++++++++++++ tfjs-backend-cpu/src/register_all_kernels.ts | 5 +- tfjs-core/src/backends/backend.ts | 6 +++ tfjs-core/src/engine.ts | 20 ++++++++ tfjs-core/src/engine_test.ts | 10 +++- 7 files changed, 143 insertions(+), 15 deletions(-) create mode 100644 tfjs-backend-cpu/src/kernels/Reshape.ts create mode 100644 tfjs-backend-cpu/src/kernels/Reshape_impl.ts diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index 8bf6be1794a..4d20503cfd3 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -59,6 +59,9 @@ export interface TensorData { // TODO(smilkov): Replace Tensor with TensorInfo when you modularize ops // that work with complex tensors. complexTensors?: {real: Tensor, imag: Tensor}; + // refCount keeps track of how many tensors reference it. Used for memory + // management. + refCount: number; } export class MathBackendCPU extends KernelBackend { @@ -91,14 +94,39 @@ export class MathBackendCPU extends KernelBackend { } } const dataId = {}; - this.data.set(dataId, {values, dtype}); + + // When creating a data bucket, we assume it is temporary, therefore its + // refCount starts from 0. When an output tensor is made, it will increase + // the refCount by 1. This makes sure the output data is not disposed. + this.data.set(dataId, {values, dtype, refCount: 0}); + return dataId; } + /** Increase refCount of a `TensorData`. */ + incRef(dataId: DataId): void { + if (this.data.has(dataId)) { + const tensorData = this.data.get(dataId); + tensorData.refCount++; + } + // Do not throw error when dataId not found for testing. Some tests + // may use the backend without actually write any data to the backend. + } + + /** Decrease refCount of a `TensorData`. */ + decRef(dataId: DataId): void { + if (this.data.has(dataId)) { + const tensorData = this.data.get(dataId); + tensorData.refCount--; + } + // Do not throw error when dataId not found for testing. Some tests + // may use the backend without actually write any data to the backend. + } + move( dataId: DataId, values: backend_util.BackendValues, shape: number[], dtype: DataType): void { - this.data.set(dataId, {values, dtype}); + this.data.set(dataId, {values, dtype, refCount: 0}); } numDataIds(): number { @@ -142,12 +170,18 @@ export class MathBackendCPU extends KernelBackend { disposeData(dataId: DataId): void { if (this.data.has(dataId)) { - const {complexTensors} = this.data.get(dataId); - if (complexTensors != null) { - complexTensors.real.dispose(); - complexTensors.imag.dispose(); + const tensorData = this.data.get(dataId); + + if (tensorData.refCount < 1) { + if (tensorData.complexTensors != null) { + // Todo(linazhao): Change to disposeData once complex, real, and imag + // kernels are modularized and real and imag becomes `TensorInfo`. + tensorData.complexTensors.real.dispose(); + tensorData.complexTensors.imag.dispose(); + } + + this.data.delete(dataId); } - this.data.delete(dataId); } } @@ -2631,10 +2665,6 @@ export class MathBackendCPU extends KernelBackend { return backend_util.castTensor(x, dtype, this); } - reshape(x: Tensor, shape: ShapeMap[R]): Tensor { - return backend_util.reshapeTensor(x, shape); - } - avgPool(x: Tensor4D, convInfo: backend_util.Conv2DInfo): Tensor4D { assertNotComplex(x, 'avgPool'); assertNotComplex(x, 'maxPool'); diff --git a/tfjs-backend-cpu/src/kernels/Reshape.ts b/tfjs-backend-cpu/src/kernels/Reshape.ts new file mode 100644 index 00000000000..ea9916072c2 --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/Reshape.ts @@ -0,0 +1,34 @@ +/** + * @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 {Reshape, ReshapeAttrs, ReshapeInputs} from '@tensorflow/tfjs-core'; +import {KernelConfig} from '@tensorflow/tfjs-core'; + +export const reshapeConfig: KernelConfig = { + kernelName: Reshape, + backendName: 'cpu', + kernelFunc: ({inputs, attrs}) => { + const {x} = inputs as ReshapeInputs; + const {shape} = attrs as {} as ReshapeAttrs; + + // Todo(linazhao): use reshapeImpl once the `TensorInfo` refCounter + // mechanism are deprecated. Right now, the engine will take care of + // increase refCount. + + return {dataId: x.dataId, shape, dtype: x.dtype}; + } +}; diff --git a/tfjs-backend-cpu/src/kernels/Reshape_impl.ts b/tfjs-backend-cpu/src/kernels/Reshape_impl.ts new file mode 100644 index 00000000000..9f0d8ad5ccd --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/Reshape_impl.ts @@ -0,0 +1,31 @@ +/** + * @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 {DataId, DataType, KernelBackend, TensorInfo} from '@tensorflow/tfjs-core'; + +// Backend kernels should use this reshape instead of the reshape kernel. + +// Internally, reshape does not create new `TensorData`, it will only increase +// the reference to the existing `TensorData`, and return a `TensorInfo` with +// the required shape. +export function reshapeImpl( + dataId: DataId, shape: number[], dtype: DataType, + backend: KernelBackend): TensorInfo { + backend.incRef(dataId); + + return {dataId, shape, dtype}; +} diff --git a/tfjs-backend-cpu/src/register_all_kernels.ts b/tfjs-backend-cpu/src/register_all_kernels.ts index a03d8fcca72..16ca2e8afb1 100644 --- a/tfjs-backend-cpu/src/register_all_kernels.ts +++ b/tfjs-backend-cpu/src/register_all_kernels.ts @@ -28,6 +28,7 @@ import {maxConfig} from './kernels/Max'; import {maxPoolWithArgmaxConfig} from './kernels/MaxPoolWithArgmax'; import {nonMaxSuppressionV4Config} from './kernels/NonMaxSuppressionV4'; import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5'; +import {reshapeConfig} from './kernels/Reshape'; import {rotateWithOffsetConfig} from './kernels/RotateWithOffset'; import {squareConfig} from './kernels/Square'; import {squaredDifferenceConfig} from './kernels/SquaredDifference'; @@ -38,8 +39,8 @@ const kernelConfigs: KernelConfig[] = [ dilation2dConfig, dilation2dBackpropInputConfig, dilation2dBackpropFilterConfig, divConfig, flipLeftRightConfig, maxPoolWithArgmaxConfig, maxConfig, nonMaxSuppressionV4Config, - nonMaxSuppressionV5Config, rotateWithOffsetConfig, squareConfig, - squaredDifferenceConfig, transposeConfig + nonMaxSuppressionV5Config, reshapeConfig, rotateWithOffsetConfig, + squareConfig, squaredDifferenceConfig, transposeConfig ]; for (const kernelConfig of kernelConfigs) { diff --git a/tfjs-core/src/backends/backend.ts b/tfjs-core/src/backends/backend.ts index 66907f2fa8d..7fed43c2764 100644 --- a/tfjs-core/src/backends/backend.ts +++ b/tfjs-core/src/backends/backend.ts @@ -113,6 +113,12 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer { write(values: BackendValues, shape: number[], dtype: DataType): DataId { return notYetImplemented('write'); } + incRef(dataId: DataId): void { + return notYetImplemented('incRef'); + } + decRef(dataId: DataId): void { + return notYetImplemented('decRef'); + } move(dataId: DataId, values: BackendValues, shape: number[], dtype: DataType): void { return notYetImplemented('move'); diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index 6d6c309df1a..84d86213f2b 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -802,7 +802,17 @@ export class Engine implements TensorTracker, DataMover { }); this.state.numBytes += bytes; } + this.state.tensorInfo.get(a.dataId).refCount++; + + // If Tensor is just created, i.e. refCount is 1, the backend will set + // the TensorData refCount during write, so no need to call incRef. + if (refCount > 1) { + this.backend.incRef(a.dataId); + } + // Todo(linazhao): Remove incRef once all the kernels use TensorData + // refCount. + if (!(a instanceof Variable)) { this.track(a); } @@ -819,6 +829,15 @@ export class Engine implements TensorTracker, DataMover { } const info = this.state.tensorInfo.get(a.dataId); const refCount = info.refCount; + + // Todo(linazhao): Move decRef into backend's disposeData method once + // refCount mechanism is completely moved from Engine to backends. + // During the migration, non-modularized kernels will still use + // TensorInfo's refCounting whereas modularized kernels will use + // TensorData's refCounting, so we need to incRef and decRef both + // TensorInfo and TensorData's refCount to keep them in sync. + this.backend.decRef(a.dataId); + if (refCount <= 1) { // Don't count bytes for complex numbers as they are counted by their // components. @@ -826,6 +845,7 @@ export class Engine implements TensorTracker, DataMover { this.state.numBytes -= info.bytes; } this.state.numDataBuffers--; + info.backend.disposeData(a.dataId); this.state.tensorInfo.delete(a.dataId); } else { diff --git a/tfjs-core/src/engine_test.ts b/tfjs-core/src/engine_test.ts index f68c0ae42df..882a63b2c22 100644 --- a/tfjs-core/src/engine_test.ts +++ b/tfjs-core/src/engine_test.ts @@ -686,7 +686,9 @@ describeWithFlags('Detects memory leaks in kernels', ALL_ENVS, () => { id: 1, dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => dataIdsCount + numDataIds: () => dataIdsCount, + incRef: (dataId: {}) => null, + decRef: (dataId: {}) => null } as TestStorage; }); @@ -712,7 +714,9 @@ describeWithFlags('Detects memory leaks in kernels', ALL_ENVS, () => { id: 1, dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => dataIdsCount + numDataIds: () => dataIdsCount, + incRef: (dataId: {}) => null, + decRef: (dataId: {}) => null } as TestStorage; }); tf.setBackend(backendName); @@ -767,6 +771,8 @@ describe('Memory allocation outside a test scope', () => { read: async (dataId: object) => storedValues, dispose: () => null, disposeData: (dataId: {}) => null, + incRef: (dataId: {}) => {}, + decRef: (dataId: {}) => {} } as TestStorage; }); tf.setBackend(backendName); From 476e19eb2f05c0ee4f02d71a9efd190a6aa0636f Mon Sep 17 00:00:00 2001 From: Na Li Date: Sun, 26 Jul 2020 17:10:57 -0700 Subject: [PATCH 02/22] Modularize kernel spacetobatchnd. --- tfjs-backend-cpu/src/backend_cpu.ts | 50 +---------- tfjs-backend-cpu/src/kernels/PadV2.ts | 66 +++++++++++++++ .../src/kernels/SpaceToBatchND.ts | 83 +++++++++++++++++++ tfjs-backend-cpu/src/register_all_kernels.ts | 6 +- 4 files changed, 154 insertions(+), 51 deletions(-) create mode 100644 tfjs-backend-cpu/src/kernels/PadV2.ts create mode 100644 tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index 4d20503cfd3..9f8d8290a4c 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -18,7 +18,7 @@ import * as tf from '@tensorflow/tfjs-core'; import {engine, env} from '@tensorflow/tfjs-core'; import {backend_util, buffer, slice_util, util} from '@tensorflow/tfjs-core'; -import {BackendTimingInfo, DataStorage, DataType, DataValues, KernelBackend, max, NumericDataType, Rank, reshape, Scalar, ShapeMap, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer, TypedArray, upcastType} from '@tensorflow/tfjs-core'; +import {BackendTimingInfo, DataStorage, DataType, DataValues, KernelBackend, max, NumericDataType, Rank, Scalar, ShapeMap, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer, TypedArray, upcastType} from '@tensorflow/tfjs-core'; import {kernel_impls} from '@tensorflow/tfjs-core'; const nonMaxSuppressionV3Impl = kernel_impls.nonMaxSuppressionV3Impl; @@ -2095,27 +2095,6 @@ export class MathBackendCPU extends KernelBackend { return tile(this.bufferSync(x), reps) as T; } - pad( - x: T, paddings: Array<[number, number]>, constantValue: number): T { - assertNotComplex(x, 'pad'); - - const outShape = paddings.map( - (p, i) => p[0] /* beforePad */ + x.shape[i] + p[1] /* afterPad */); - const start = paddings.map(p => p[0]); - const xBuffer = this.bufferSync(x); - const buffer = tf.buffer(outShape, x.dtype as 'float32'); - if (constantValue !== 0) { - buffer.values.fill(constantValue); - } - - for (let i = 0; i < x.size; i++) { - const coords = xBuffer.indexToLoc(i); - const outCoords = coords.map((c, i) => c + start[i]); - buffer.set(xBuffer.get(...coords), ...outCoords); - } - return buffer.toTensor() as T; - } - gather(x: T, indices: Tensor1D, axis: number): T { assertNotComplex([x, indices], 'gather'); @@ -2158,33 +2137,6 @@ export class MathBackendCPU extends KernelBackend { .slice(sliceBeginCoords, sliceSize) as T; } - spaceToBatchND( - x: T, blockShape: number[], paddings: Array<[number, number]>): T { - assertNotComplex([x], 'spaceToBatchND'); - - const prod = blockShape.reduce((a, b) => a * b); - - const completePaddings: Array<[number, number]> = [[0, 0]]; - completePaddings.push(...paddings); - for (let i = 1 + blockShape.length; i < x.shape.length; ++i) { - completePaddings.push([0, 0]); - } - - const paddedX = x.pad(completePaddings); - - const reshapedPaddedShape = - backend_util.getReshaped(paddedX.shape, blockShape, prod, false); - const permutedReshapedPaddedPermutation = backend_util.getPermuted( - reshapedPaddedShape.length, blockShape.length, false); - const flattenShape = backend_util.getReshapedPermuted( - paddedX.shape, blockShape, prod, false); - - const paddedXT = tf.transpose( - paddedX.reshape(reshapedPaddedShape), - permutedReshapedPaddedPermutation); - return reshape(paddedXT, flattenShape) as T; - } - maxPool(x: Tensor4D, convInfo: backend_util.Conv2DInfo): Tensor4D { assertNotComplex(x, 'maxPool'); const xValues = this.readSync(x.dataId) as TypedArray; diff --git a/tfjs-backend-cpu/src/kernels/PadV2.ts b/tfjs-backend-cpu/src/kernels/PadV2.ts new file mode 100644 index 00000000000..bbbf94fc6ca --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/PadV2.ts @@ -0,0 +1,66 @@ +/** + * @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 {NumericDataType, PadV2, PadV2Attrs, PadV2Inputs, TypedArray, util} from '@tensorflow/tfjs-core'; +import {KernelConfig} from '@tensorflow/tfjs-core'; + +import {MathBackendCPU} from '../backend_cpu'; +import {assertNotComplex} from '../cpu_util'; + +export const padV2Config: KernelConfig = { + kernelName: PadV2, + backendName: 'cpu', + kernelFunc: ({inputs, backend, attrs}) => { + const {x} = inputs as PadV2Inputs; + const {paddings, constantValue} = attrs as {} as PadV2Attrs; + const cpuBackend = backend as MathBackendCPU; + + assertNotComplex(x, 'pad'); + + const outShape = paddings.map( + (p, i) => p[0] /* beforePad */ + x.shape[i] + p[1] /* afterPad */); + + const start = paddings.map(p => p[0]); + + const xVals = cpuBackend.data.get(x.dataId).values as TypedArray; + const xSize = util.sizeFromShape(x.shape); + const xRank = x.shape.length; + const xStrides = util.computeStrides(x.shape); + + const resultSize = util.sizeFromShape(outShape); + const resultRank = outShape.length; + const resultStrides = util.computeStrides(outShape); + const resVals = + util.getTypedArrayFromDType(x.dtype as NumericDataType, resultSize); + + if (constantValue !== 0) { + resVals.fill(constantValue); + } + + for (let i = 0; i < xSize; i++) { + const coords = util.indexToLoc(i, xRank, xStrides); + const outCoords = coords.map((c, i) => c + start[i]); + const outIndex = util.locToIndex(outCoords, resultRank, resultStrides); + + resVals[outIndex] = xVals[i]; + } + + const outId = cpuBackend.write(resVals, outShape, x.dtype); + + return {dataId: outId, shape: outShape, dtype: x.dtype}; + } +}; diff --git a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts new file mode 100644 index 00000000000..305606d8659 --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts @@ -0,0 +1,83 @@ +/** + * @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 {backend_util, SpaceToBatchND, SpaceToBatchNDAttrs, SpaceToBatchNDInputs, TensorInfo} from '@tensorflow/tfjs-core'; +import {KernelConfig} from '@tensorflow/tfjs-core'; + +import {MathBackendCPU} from '../backend_cpu'; +import {assertNotComplex} from '../cpu_util'; + +import {padV2Config} from './PadV2'; +import {reshapeConfig} from './Reshape'; +import {transposeConfig} from './Transpose'; + +export const spaceToBatchNDConfig: KernelConfig = { + kernelName: SpaceToBatchND, + backendName: 'cpu', + kernelFunc: ({inputs, backend, attrs}) => { + const {x} = inputs as SpaceToBatchNDInputs; + const {blockShape, paddings} = attrs as {} as SpaceToBatchNDAttrs; + const cpuBackend = backend as MathBackendCPU; + + assertNotComplex([x], 'spaceToBatchND'); + + const prod = blockShape.reduce((a, b) => a * b); + + const completePaddings: Array<[number, number]> = [[0, 0]]; + completePaddings.push(...(paddings as Array<[number, number]>)); + + for (let i = 1 + blockShape.length; i < x.shape.length; ++i) { + completePaddings.push([0, 0]); + } + + const paddedX = padV2Config.kernelFunc({ + inputs: {x}, + backend, + attrs: {completePaddings, constantValue: 0} + }) as TensorInfo; + + const reshapedPaddedShape = + backend_util.getReshaped(paddedX.shape, blockShape, prod, false); + + const permutedReshapedPaddedPermutation = backend_util.getPermuted( + reshapedPaddedShape.length, blockShape.length, false); + + const flattenShape = backend_util.getReshapedPermuted( + paddedX.shape, blockShape, prod, false); + + const paddedXReshaped = reshapeConfig.kernelFunc({ + inputs: {x: paddedX}, + backend, + attrs: {shape: reshapedPaddedShape} + }) as TensorInfo; + + const paddedXT = transposeConfig.kernelFunc({ + inputs: {x: paddedXReshaped}, + backend, + attrs: {perm: permutedReshapedPaddedPermutation} + }) as TensorInfo; + + const result = reshapeConfig.kernelFunc( + {inputs: {x: paddedXT}, backend, attrs: {shape: flattenShape}}); + + cpuBackend.disposeData(paddedX.dataId); + cpuBackend.disposeData(paddedXReshaped.dataId); + cpuBackend.disposeData(paddedXT.dataId); + + return result; + } +}; diff --git a/tfjs-backend-cpu/src/register_all_kernels.ts b/tfjs-backend-cpu/src/register_all_kernels.ts index 16ca2e8afb1..70ec59226df 100644 --- a/tfjs-backend-cpu/src/register_all_kernels.ts +++ b/tfjs-backend-cpu/src/register_all_kernels.ts @@ -28,8 +28,10 @@ import {maxConfig} from './kernels/Max'; import {maxPoolWithArgmaxConfig} from './kernels/MaxPoolWithArgmax'; import {nonMaxSuppressionV4Config} from './kernels/NonMaxSuppressionV4'; import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5'; +import {padV2Config} from './kernels/PadV2'; import {reshapeConfig} from './kernels/Reshape'; import {rotateWithOffsetConfig} from './kernels/RotateWithOffset'; +import {spaceToBatchNDConfig} from './kernels/SpaceToBatchND'; import {squareConfig} from './kernels/Square'; import {squaredDifferenceConfig} from './kernels/SquaredDifference'; import {transposeConfig} from './kernels/Transpose'; @@ -39,8 +41,8 @@ const kernelConfigs: KernelConfig[] = [ dilation2dConfig, dilation2dBackpropInputConfig, dilation2dBackpropFilterConfig, divConfig, flipLeftRightConfig, maxPoolWithArgmaxConfig, maxConfig, nonMaxSuppressionV4Config, - nonMaxSuppressionV5Config, reshapeConfig, rotateWithOffsetConfig, - squareConfig, squaredDifferenceConfig, transposeConfig + nonMaxSuppressionV5Config, padV2Config, reshapeConfig, rotateWithOffsetConfig, + spaceToBatchNDConfig, squareConfig, squaredDifferenceConfig, transposeConfig ]; for (const kernelConfig of kernelConfigs) { From b18d2105f6e205fc6159ebfbf5d4ca4cc0ded568 Mon Sep 17 00:00:00 2001 From: Na Li Date: Sun, 26 Jul 2020 18:53:55 -0700 Subject: [PATCH 03/22] Modularize kernel spacetobatchnd. --- tfjs-backend-cpu/src/backend_cpu.ts | 9 ++++----- tfjs-backend-cpu/src/kernels/Reshape.ts | 8 ++++---- .../src/kernels/SpaceToBatchND.ts | 2 +- tfjs-core/src/engine.ts | 20 ++++++------------- tfjs-core/src/kernel_registry_test.ts | 12 ++++++++--- 5 files changed, 24 insertions(+), 27 deletions(-) diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index 9f8d8290a4c..87bfb4409d9 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -95,10 +95,7 @@ export class MathBackendCPU extends KernelBackend { } const dataId = {}; - // When creating a data bucket, we assume it is temporary, therefore its - // refCount starts from 0. When an output tensor is made, it will increase - // the refCount by 1. This makes sure the output data is not disposed. - this.data.set(dataId, {values, dtype, refCount: 0}); + this.data.set(dataId, {values, dtype, refCount: 1}); return dataId; } @@ -126,7 +123,7 @@ export class MathBackendCPU extends KernelBackend { move( dataId: DataId, values: backend_util.BackendValues, shape: number[], dtype: DataType): void { - this.data.set(dataId, {values, dtype, refCount: 0}); + this.data.set(dataId, {values, dtype, refCount: 1}); } numDataIds(): number { @@ -172,6 +169,8 @@ export class MathBackendCPU extends KernelBackend { if (this.data.has(dataId)) { const tensorData = this.data.get(dataId); + tensorData.refCount--; + if (tensorData.refCount < 1) { if (tensorData.complexTensors != null) { // Todo(linazhao): Change to disposeData once complex, real, and imag diff --git a/tfjs-backend-cpu/src/kernels/Reshape.ts b/tfjs-backend-cpu/src/kernels/Reshape.ts index ea9916072c2..452f6990a5a 100644 --- a/tfjs-backend-cpu/src/kernels/Reshape.ts +++ b/tfjs-backend-cpu/src/kernels/Reshape.ts @@ -17,17 +17,17 @@ import {Reshape, ReshapeAttrs, ReshapeInputs} from '@tensorflow/tfjs-core'; import {KernelConfig} from '@tensorflow/tfjs-core'; +import {MathBackendCPU} from '../backend_cpu'; export const reshapeConfig: KernelConfig = { kernelName: Reshape, backendName: 'cpu', - kernelFunc: ({inputs, attrs}) => { + kernelFunc: ({inputs, backend, attrs}) => { const {x} = inputs as ReshapeInputs; const {shape} = attrs as {} as ReshapeAttrs; + const cpuBackend = backend as MathBackendCPU; - // Todo(linazhao): use reshapeImpl once the `TensorInfo` refCounter - // mechanism are deprecated. Right now, the engine will take care of - // increase refCount. + cpuBackend.incRef(x.dataId); return {dataId: x.dataId, shape, dtype: x.dtype}; } diff --git a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts index 305606d8659..3e6d2f072a0 100644 --- a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts +++ b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts @@ -47,7 +47,7 @@ export const spaceToBatchNDConfig: KernelConfig = { const paddedX = padV2Config.kernelFunc({ inputs: {x}, backend, - attrs: {completePaddings, constantValue: 0} + attrs: {paddings: completePaddings, constantValue: 0} }) as TensorInfo; const reshapedPaddedShape = diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index 84d86213f2b..4e02e63368b 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -805,14 +805,6 @@ export class Engine implements TensorTracker, DataMover { this.state.tensorInfo.get(a.dataId).refCount++; - // If Tensor is just created, i.e. refCount is 1, the backend will set - // the TensorData refCount during write, so no need to call incRef. - if (refCount > 1) { - this.backend.incRef(a.dataId); - } - // Todo(linazhao): Remove incRef once all the kernels use TensorData - // refCount. - if (!(a instanceof Variable)) { this.track(a); } @@ -830,13 +822,13 @@ export class Engine implements TensorTracker, DataMover { const info = this.state.tensorInfo.get(a.dataId); const refCount = info.refCount; - // Todo(linazhao): Move decRef into backend's disposeData method once - // refCount mechanism is completely moved from Engine to backends. - // During the migration, non-modularized kernels will still use - // TensorInfo's refCounting whereas modularized kernels will use + // Todo(linazhao): Move decRef into backend's disposeData method once + // refCount mechanism is completely moved from Engine to backends. + // During the migration, non-modularized kernels will still use + // TensorInfo's refCounting whereas modularized kernels will use // TensorData's refCounting, so we need to incRef and decRef both - // TensorInfo and TensorData's refCount to keep them in sync. - this.backend.decRef(a.dataId); + // TensorInfo and TensorData's refCount to keep them in sync. + info.backend.decRef(a.dataId); if (refCount <= 1) { // Don't count bytes for complex numbers as they are counted by their diff --git a/tfjs-core/src/kernel_registry_test.ts b/tfjs-core/src/kernel_registry_test.ts index 352fce14391..f799a2e5ee3 100644 --- a/tfjs-core/src/kernel_registry_test.ts +++ b/tfjs-core/src/kernel_registry_test.ts @@ -100,7 +100,9 @@ describeWithFlags('kernel_registry', ALL_ENVS, () => { id: 1, dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => 0 + numDataIds: () => 0, + incRef: (dataId: {}) => null, + decRef: (dataId: {}) => null } as TestBackend; }); tf.registerBackend('backend2', () => { @@ -108,7 +110,9 @@ describeWithFlags('kernel_registry', ALL_ENVS, () => { id: 2, dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => 0 + numDataIds: () => 0, + incRef: (dataId: {}) => null, + decRef: (dataId: {}) => null } as TestBackend; }); @@ -148,7 +152,9 @@ describeWithFlags('kernel_registry', ALL_ENVS, () => { const customBackend = { dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => 0 + numDataIds: () => 0, + incRef: (dataId: {}) => null, + decRef: (dataId: {}) => null } as TestBackend; tf.registerBackend(backendName, () => customBackend); From 6f288b204a7563e0dbd521ec5a7d543ca5ca4dcf Mon Sep 17 00:00:00 2001 From: Na Li Date: Sun, 26 Jul 2020 18:56:50 -0700 Subject: [PATCH 04/22] Modularize kernel spacetobatchnd. --- tfjs-backend-cpu/src/kernels/Reshape_impl.ts | 31 -------------------- tfjs-core/src/engine.ts | 2 +- 2 files changed, 1 insertion(+), 32 deletions(-) delete mode 100644 tfjs-backend-cpu/src/kernels/Reshape_impl.ts diff --git a/tfjs-backend-cpu/src/kernels/Reshape_impl.ts b/tfjs-backend-cpu/src/kernels/Reshape_impl.ts deleted file mode 100644 index 9f0d8ad5ccd..00000000000 --- a/tfjs-backend-cpu/src/kernels/Reshape_impl.ts +++ /dev/null @@ -1,31 +0,0 @@ -/** - * @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 {DataId, DataType, KernelBackend, TensorInfo} from '@tensorflow/tfjs-core'; - -// Backend kernels should use this reshape instead of the reshape kernel. - -// Internally, reshape does not create new `TensorData`, it will only increase -// the reference to the existing `TensorData`, and return a `TensorInfo` with -// the required shape. -export function reshapeImpl( - dataId: DataId, shape: number[], dtype: DataType, - backend: KernelBackend): TensorInfo { - backend.incRef(dataId); - - return {dataId, shape, dtype}; -} diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index 4e02e63368b..bcc2d7e0993 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -826,7 +826,7 @@ export class Engine implements TensorTracker, DataMover { // refCount mechanism is completely moved from Engine to backends. // During the migration, non-modularized kernels will still use // TensorInfo's refCounting whereas modularized kernels will use - // TensorData's refCounting, so we need to incRef and decRef both + // TensorData's refCounting, so we need to decRef both // TensorInfo and TensorData's refCount to keep them in sync. info.backend.decRef(a.dataId); From aa20f7ea32abb2fa6317335a41954a087ed27fc4 Mon Sep 17 00:00:00 2001 From: Na Li Date: Wed, 29 Jul 2020 11:42:50 -0700 Subject: [PATCH 05/22] Move decRef to the right place. --- tfjs-backend-cpu/src/backend_cpu_test.ts | 42 +++++++++++++++++------- tfjs-core/src/engine.ts | 16 ++++----- 2 files changed, 39 insertions(+), 19 deletions(-) diff --git a/tfjs-backend-cpu/src/backend_cpu_test.ts b/tfjs-backend-cpu/src/backend_cpu_test.ts index 768ffe578db..d44854f5a19 100644 --- a/tfjs-backend-cpu/src/backend_cpu_test.ts +++ b/tfjs-backend-cpu/src/backend_cpu_test.ts @@ -16,7 +16,8 @@ */ import * as tf from '@tensorflow/tfjs-core'; -import {engine, test_util, util} from '@tensorflow/tfjs-core'; +import {engine, Tensor, test_util, util} from '@tensorflow/tfjs-core'; + const {expectArraysClose, expectArraysEqual} = test_util; // tslint:disable-next-line: no-imports-from-dist import {describeWithFlags, ALL_ENVS} from '@tensorflow/tfjs-core/dist/jasmine_util'; @@ -131,19 +132,38 @@ describeWithFlags('memory cpu', ALL_ENVS, () => { '(2 bytes per character)'; expect(mem.reasons.indexOf(expectedReasonString) >= 0).toBe(true); }); -}); -describeWithFlags('CPU backend has sync init', ALL_ENVS, () => { - it('can do matmul without waiting for ready', async () => { + it('does not have memory leak.', async () => { tf.registerBackend('my-cpu', () => { return new MathBackendCPU(); }); - tf.setBackend('my-cpu'); - const a = tf.tensor1d([5]); - const b = tf.tensor1d([3]); - const res = tf.dot(a, b); - expectArraysClose(await res.data(), 15); - tf.dispose([a, b, res]); - tf.removeBackend('my-cpu'); + await tf.setBackend('my-cpu'); + + tf.registerKernel({ + kernelName: 'SimpleKernel', + backendName: 'my-cpu', + kernelFunc: ({backend}) => { + const cpuBackend = backend as MathBackendCPU; + const outId = cpuBackend.write(new Float32Array(1), [], 'float32'); + return {dtype: 'float32', shape: [], dataId: outId}; + } + }); + + const beforeDataIds = tf.engine().backend.numDataIds(); + + const res = tf.engine().runKernel('SimpleKernel', {}, {}) as Tensor; + + expectArraysClose(await res.data(), [0]); + expectArraysEqual(res.shape, []); + + const afterResDataIds = tf.engine().backend.numDataIds(); + expect(afterResDataIds).toEqual(beforeDataIds + 1); + + res.dispose(); + + const afterDisposeDataIds = tf.engine().backend.numDataIds(); + expect(afterDisposeDataIds).toEqual(beforeDataIds); + + tf.unregisterKernel('SimpleKernel', tf.getBackend()); }); }); diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index bcc2d7e0993..819b9db002f 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -822,14 +822,6 @@ export class Engine implements TensorTracker, DataMover { const info = this.state.tensorInfo.get(a.dataId); const refCount = info.refCount; - // Todo(linazhao): Move decRef into backend's disposeData method once - // refCount mechanism is completely moved from Engine to backends. - // During the migration, non-modularized kernels will still use - // TensorInfo's refCounting whereas modularized kernels will use - // TensorData's refCounting, so we need to decRef both - // TensorInfo and TensorData's refCount to keep them in sync. - info.backend.decRef(a.dataId); - if (refCount <= 1) { // Don't count bytes for complex numbers as they are counted by their // components. @@ -842,6 +834,14 @@ export class Engine implements TensorTracker, DataMover { this.state.tensorInfo.delete(a.dataId); } else { this.state.tensorInfo.get(a.dataId).refCount--; + + // Todo(linazhao): Move decRef into backend's disposeData method once + // refCount mechanism is completely moved from Engine to backends. + // During the migration, non-modularized kernels will still use + // TensorInfo's refCounting whereas modularized kernels will use + // TensorData's refCounting, so we need to decRef both + // TensorInfo and TensorData's refCount to keep them in sync. + info.backend.decRef(a.dataId); } // TODO(nsthorat): Construct an error and save the stack trace for // debugging when in debug mode. Creating a stack trace is too expensive From 0342781c1002201c48ae39a4531e8bca933a39af Mon Sep 17 00:00:00 2001 From: Na Li Date: Wed, 29 Jul 2020 16:13:40 -0700 Subject: [PATCH 06/22] Add memory test. --- tfjs-backend-cpu/src/backend_cpu_test.ts | 28 ++++++--------------- tfjs-core/src/ops/space_to_batch_nd_test.ts | 13 ++++++++++ 2 files changed, 20 insertions(+), 21 deletions(-) diff --git a/tfjs-backend-cpu/src/backend_cpu_test.ts b/tfjs-backend-cpu/src/backend_cpu_test.ts index d44854f5a19..25ed315ee6a 100644 --- a/tfjs-backend-cpu/src/backend_cpu_test.ts +++ b/tfjs-backend-cpu/src/backend_cpu_test.ts @@ -133,37 +133,23 @@ describeWithFlags('memory cpu', ALL_ENVS, () => { expect(mem.reasons.indexOf(expectedReasonString) >= 0).toBe(true); }); - it('does not have memory leak.', async () => { - tf.registerBackend('my-cpu', () => { - return new MathBackendCPU(); - }); - await tf.setBackend('my-cpu'); - - tf.registerKernel({ - kernelName: 'SimpleKernel', - backendName: 'my-cpu', - kernelFunc: ({backend}) => { - const cpuBackend = backend as MathBackendCPU; - const outId = cpuBackend.write(new Float32Array(1), [], 'float32'); - return {dtype: 'float32', shape: [], dataId: outId}; - } - }); - + it('does not have memory leak with reshape.', async () => { const beforeDataIds = tf.engine().backend.numDataIds(); - const res = tf.engine().runKernel('SimpleKernel', {}, {}) as Tensor; + const x = tf.tensor1d([1, 1, 1, 1]); + const res = + tf.engine().runKernel('Reshape', {x}, {shape: [2, 2]}) as Tensor; - expectArraysClose(await res.data(), [0]); - expectArraysEqual(res.shape, []); + expectArraysClose(await res.data(), [1, 1, 1, 1]); + expectArraysEqual(res.shape, [2, 2]); const afterResDataIds = tf.engine().backend.numDataIds(); expect(afterResDataIds).toEqual(beforeDataIds + 1); + x.dispose(); res.dispose(); const afterDisposeDataIds = tf.engine().backend.numDataIds(); expect(afterDisposeDataIds).toEqual(beforeDataIds); - - tf.unregisterKernel('SimpleKernel', tf.getBackend()); }); }); diff --git a/tfjs-core/src/ops/space_to_batch_nd_test.ts b/tfjs-core/src/ops/space_to_batch_nd_test.ts index 11da95ac170..f6485ec8f25 100644 --- a/tfjs-core/src/ops/space_to_batch_nd_test.ts +++ b/tfjs-core/src/ops/space_to_batch_nd_test.ts @@ -21,6 +21,8 @@ import {expectArraysClose} from '../test_util'; describeWithFlags('spaceToBatchND', ALL_ENVS, () => { it('tensor4d, input shape=[1, 2, 2, 1], blockShape=[2, 2]', async () => { + const initialDataIds = tf.engine().backend.numDataIds(); + const t = tf.tensor4d([[[[1], [2]], [[3], [4]]]], [1, 2, 2, 1]); const blockShape = [2, 2]; const paddings = [[0, 0], [0, 0]]; @@ -28,6 +30,17 @@ describeWithFlags('spaceToBatchND', ALL_ENVS, () => { const res = tf.spaceToBatchND(t, blockShape, paddings); expect(res.shape).toEqual([4, 1, 1, 1]); expectArraysClose(await res.data(), [1, 2, 3, 4]); + + const afterResDataIds = tf.engine().backend.numDataIds(); + // 1 input tensor and 1 result tensor + expect(afterResDataIds).toEqual(initialDataIds + 2); + + t.dispose(); + res.dispose(); + + const afterDisposeDataIds = tf.engine().backend.numDataIds(); + + expect(afterDisposeDataIds).toEqual(initialDataIds); }); it('tensor4d, input shape=[1, 2, 2, 3], blockShape=[2, 2]', async () => { From 19ad1ddfa5a982e3fc2e4c27a1be76d192987910 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 30 Jul 2020 19:53:27 -0700 Subject: [PATCH 07/22] Clean up. --- tfjs-backend-cpu/src/backend_cpu.ts | 10 -- tfjs-backend-cpu/src/kernels/PadV2.ts | 71 ++++++----- tfjs-backend-cpu/src/kernels/Reshape.ts | 27 ++-- tfjs-backend-cpu/src/kernels/Reshape_test.ts | 45 +++++++ .../src/kernels/SpaceToBatchND.ts | 117 ++++++++++-------- .../src/kernels/Transpose_impl.ts | 4 +- tfjs-backend-cpu/src/utils/kernel_utils.ts | 6 +- tfjs-backend-wasm/src/backend_wasm.ts | 2 + tfjs-backend-webgl/src/backend_webgl.ts | 2 + tfjs-backend-webgpu/src/backend_webgpu.ts | 2 + tfjs-core/src/backends/backend.ts | 3 - tfjs-core/src/engine.ts | 10 +- tfjs-core/src/util.ts | 12 +- tfjs-node/src/nodejs_kernel_backend.ts | 2 + 14 files changed, 187 insertions(+), 126 deletions(-) create mode 100644 tfjs-backend-cpu/src/kernels/Reshape_test.ts diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index 87bfb4409d9..a160a95ffb2 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -100,16 +100,6 @@ export class MathBackendCPU extends KernelBackend { return dataId; } - /** Increase refCount of a `TensorData`. */ - incRef(dataId: DataId): void { - if (this.data.has(dataId)) { - const tensorData = this.data.get(dataId); - tensorData.refCount++; - } - // Do not throw error when dataId not found for testing. Some tests - // may use the backend without actually write any data to the backend. - } - /** Decrease refCount of a `TensorData`. */ decRef(dataId: DataId): void { if (this.data.has(dataId)) { diff --git a/tfjs-backend-cpu/src/kernels/PadV2.ts b/tfjs-backend-cpu/src/kernels/PadV2.ts index bbbf94fc6ca..df278551e37 100644 --- a/tfjs-backend-cpu/src/kernels/PadV2.ts +++ b/tfjs-backend-cpu/src/kernels/PadV2.ts @@ -15,52 +15,55 @@ * ============================================================================= */ -import {NumericDataType, PadV2, PadV2Attrs, PadV2Inputs, TypedArray, util} from '@tensorflow/tfjs-core'; -import {KernelConfig} from '@tensorflow/tfjs-core'; +import {KernelConfig, KernelFunc, NumericDataType, PadV2, PadV2Attrs, PadV2Inputs, TensorInfo, TypedArray, util} from '@tensorflow/tfjs-core'; import {MathBackendCPU} from '../backend_cpu'; import {assertNotComplex} from '../cpu_util'; -export const padV2Config: KernelConfig = { - kernelName: PadV2, - backendName: 'cpu', - kernelFunc: ({inputs, backend, attrs}) => { - const {x} = inputs as PadV2Inputs; - const {paddings, constantValue} = attrs as {} as PadV2Attrs; - const cpuBackend = backend as MathBackendCPU; +function padV2( + args: {inputs: PadV2Inputs, backend: MathBackendCPU, attrs: PadV2Attrs}): + TensorInfo { + const {inputs, backend, attrs} = args; + const {x} = inputs; + const {paddings, constantValue} = attrs; - assertNotComplex(x, 'pad'); + assertNotComplex(x, 'pad'); - const outShape = paddings.map( - (p, i) => p[0] /* beforePad */ + x.shape[i] + p[1] /* afterPad */); + const outShape = paddings.map( + (p, i) => p[0] /* beforePad */ + x.shape[i] + p[1] /* afterPad */); - const start = paddings.map(p => p[0]); + const start = paddings.map(p => p[0]); - const xVals = cpuBackend.data.get(x.dataId).values as TypedArray; - const xSize = util.sizeFromShape(x.shape); - const xRank = x.shape.length; - const xStrides = util.computeStrides(x.shape); + const xVals = backend.data.get(x.dataId).values as TypedArray; + const xSize = util.sizeFromShape(x.shape); + const xRank = x.shape.length; + const xStrides = util.computeStrides(x.shape); - const resultSize = util.sizeFromShape(outShape); - const resultRank = outShape.length; - const resultStrides = util.computeStrides(outShape); - const resVals = - util.getTypedArrayFromDType(x.dtype as NumericDataType, resultSize); + const resultSize = util.sizeFromShape(outShape); + const resultRank = outShape.length; + const resultStrides = util.computeStrides(outShape); + const resVals = + util.getTypedArrayFromDType(x.dtype as NumericDataType, resultSize); + + if (constantValue !== 0) { + resVals.fill(constantValue); + } - if (constantValue !== 0) { - resVals.fill(constantValue); - } + for (let i = 0; i < xSize; i++) { + const coords = util.indexToCoord(i, xRank, xStrides); + const outCoords = coords.map((c, i) => c + start[i]); + const outIndex = util.coordToIndex(outCoords, resultRank, resultStrides); - for (let i = 0; i < xSize; i++) { - const coords = util.indexToLoc(i, xRank, xStrides); - const outCoords = coords.map((c, i) => c + start[i]); - const outIndex = util.locToIndex(outCoords, resultRank, resultStrides); + resVals[outIndex] = xVals[i]; + } - resVals[outIndex] = xVals[i]; - } + const outId = backend.write(resVals, outShape, x.dtype); - const outId = cpuBackend.write(resVals, outShape, x.dtype); + return {dataId: outId, shape: outShape, dtype: x.dtype}; +} - return {dataId: outId, shape: outShape, dtype: x.dtype}; - } +export const padV2Config: KernelConfig = { + kernelName: PadV2, + backendName: 'cpu', + kernelFunc: padV2 as {} as KernelFunc }; diff --git a/tfjs-backend-cpu/src/kernels/Reshape.ts b/tfjs-backend-cpu/src/kernels/Reshape.ts index 452f6990a5a..968af215399 100644 --- a/tfjs-backend-cpu/src/kernels/Reshape.ts +++ b/tfjs-backend-cpu/src/kernels/Reshape.ts @@ -15,20 +15,25 @@ * ============================================================================= */ -import {Reshape, ReshapeAttrs, ReshapeInputs} from '@tensorflow/tfjs-core'; -import {KernelConfig} from '@tensorflow/tfjs-core'; +import {KernelConfig, KernelFunc, Reshape, ReshapeAttrs, ReshapeInputs} from '@tensorflow/tfjs-core'; + import {MathBackendCPU} from '../backend_cpu'; +function reshape( + args: + {inputs: ReshapeInputs, backend: MathBackendCPU, attrs: ReshapeAttrs}) { + const {inputs, backend, attrs} = args; + const {x} = inputs; + const {shape} = attrs; + const cpuBackend = backend; + + cpuBackend.incRef(x.dataId); + + return {dataId: x.dataId, shape, dtype: x.dtype}; +} + export const reshapeConfig: KernelConfig = { kernelName: Reshape, backendName: 'cpu', - kernelFunc: ({inputs, backend, attrs}) => { - const {x} = inputs as ReshapeInputs; - const {shape} = attrs as {} as ReshapeAttrs; - const cpuBackend = backend as MathBackendCPU; - - cpuBackend.incRef(x.dataId); - - return {dataId: x.dataId, shape, dtype: x.dtype}; - } + kernelFunc: reshape as {} as KernelFunc }; diff --git a/tfjs-backend-cpu/src/kernels/Reshape_test.ts b/tfjs-backend-cpu/src/kernels/Reshape_test.ts new file mode 100644 index 00000000000..0d7903adea2 --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/Reshape_test.ts @@ -0,0 +1,45 @@ +/** + * @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 * as tf from '@tensorflow/tfjs-core'; +import {Tensor, test_util} from '@tensorflow/tfjs-core'; + +const {expectArraysClose, expectArraysEqual} = test_util; +// tslint:disable-next-line: no-imports-from-dist +import {describeWithFlags, ALL_ENVS} from '@tensorflow/tfjs-core/dist/jasmine_util'; + +describeWithFlags('Reshape.', ALL_ENVS, () => { + it('does not have memory leak.', async () => { + const beforeDataIds = tf.engine().backend.numDataIds(); + + const x = tf.tensor1d([1, 1, 1, 1]); + const res = + tf.engine().runKernel('Reshape', {x}, {shape: [2, 2]}) as Tensor; + + expectArraysClose(await res.data(), [1, 1, 1, 1]); + expectArraysEqual(res.shape, [2, 2]); + + const afterResDataIds = tf.engine().backend.numDataIds(); + expect(afterResDataIds).toEqual(beforeDataIds + 1); + + x.dispose(); + res.dispose(); + + const afterDisposeDataIds = tf.engine().backend.numDataIds(); + expect(afterDisposeDataIds).toEqual(beforeDataIds); + }); +}); diff --git a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts index 3e6d2f072a0..4fa4eb3e4b0 100644 --- a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts +++ b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts @@ -15,8 +15,7 @@ * ============================================================================= */ -import {backend_util, SpaceToBatchND, SpaceToBatchNDAttrs, SpaceToBatchNDInputs, TensorInfo} from '@tensorflow/tfjs-core'; -import {KernelConfig} from '@tensorflow/tfjs-core'; +import {backend_util, KernelConfig, KernelFunc, NamedAttrMap, ReshapeAttrs, ReshapeInputs, SpaceToBatchND, SpaceToBatchNDAttrs, SpaceToBatchNDInputs, TensorInfo, TransposeAttrs, TransposeInputs} from '@tensorflow/tfjs-core'; import {MathBackendCPU} from '../backend_cpu'; import {assertNotComplex} from '../cpu_util'; @@ -25,59 +24,75 @@ import {padV2Config} from './PadV2'; import {reshapeConfig} from './Reshape'; import {transposeConfig} from './Transpose'; -export const spaceToBatchNDConfig: KernelConfig = { - kernelName: SpaceToBatchND, - backendName: 'cpu', - kernelFunc: ({inputs, backend, attrs}) => { - const {x} = inputs as SpaceToBatchNDInputs; - const {blockShape, paddings} = attrs as {} as SpaceToBatchNDAttrs; - const cpuBackend = backend as MathBackendCPU; - - assertNotComplex([x], 'spaceToBatchND'); - - const prod = blockShape.reduce((a, b) => a * b); - - const completePaddings: Array<[number, number]> = [[0, 0]]; - completePaddings.push(...(paddings as Array<[number, number]>)); - - for (let i = 1 + blockShape.length; i < x.shape.length; ++i) { - completePaddings.push([0, 0]); - } - - const paddedX = padV2Config.kernelFunc({ - inputs: {x}, - backend, - attrs: {paddings: completePaddings, constantValue: 0} - }) as TensorInfo; +function spaceToBatchND(args: { + inputs: SpaceToBatchNDInputs, + backend: MathBackendCPU, + attrs: SpaceToBatchNDAttrs +}): TensorInfo { + const {inputs, backend, attrs} = args; + const {x} = inputs; + const {blockShape, paddings} = attrs; - const reshapedPaddedShape = - backend_util.getReshaped(paddedX.shape, blockShape, prod, false); + assertNotComplex([x], 'spaceToBatchND'); - const permutedReshapedPaddedPermutation = backend_util.getPermuted( - reshapedPaddedShape.length, blockShape.length, false); + const prod = blockShape.reduce((a, b) => a * b); - const flattenShape = backend_util.getReshapedPermuted( - paddedX.shape, blockShape, prod, false); + const completePaddings: Array<[number, number]> = [[0, 0]]; + completePaddings.push(...(paddings as Array<[number, number]>)); - const paddedXReshaped = reshapeConfig.kernelFunc({ - inputs: {x: paddedX}, - backend, - attrs: {shape: reshapedPaddedShape} - }) as TensorInfo; - - const paddedXT = transposeConfig.kernelFunc({ - inputs: {x: paddedXReshaped}, - backend, - attrs: {perm: permutedReshapedPaddedPermutation} - }) as TensorInfo; - - const result = reshapeConfig.kernelFunc( - {inputs: {x: paddedXT}, backend, attrs: {shape: flattenShape}}); + for (let i = 1 + blockShape.length; i < x.shape.length; ++i) { + completePaddings.push([0, 0]); + } - cpuBackend.disposeData(paddedX.dataId); - cpuBackend.disposeData(paddedXReshaped.dataId); - cpuBackend.disposeData(paddedXT.dataId); + const paddedX = padV2Config.kernelFunc({ + inputs: {x}, + backend, + attrs: {paddings: completePaddings, constantValue: 0} + }) as TensorInfo; + + const reshapedPaddedShape = + backend_util.getReshaped(paddedX.shape, blockShape, prod, false); + + const permutedReshapedPaddedPermutation = backend_util.getPermuted( + reshapedPaddedShape.length, blockShape.length, false); + + const flattenShape = + backend_util.getReshapedPermuted(paddedX.shape, blockShape, prod, false); + + const reshapeInputs: ReshapeInputs = {x: paddedX}; + const reshapeAttrs: ReshapeAttrs = {shape: reshapedPaddedShape}; + const paddedXReshaped = reshapeConfig.kernelFunc({ + inputs: reshapeInputs, + backend, + attrs: reshapeAttrs as {} as NamedAttrMap + }) as TensorInfo; + + const transposeInputs: TransposeInputs = {x: paddedXReshaped}; + const transposeAttrs: + TransposeAttrs = {perm: permutedReshapedPaddedPermutation}; + const paddedXT = transposeConfig.kernelFunc({ + inputs: transposeInputs, + backend, + attrs: transposeAttrs as {} as NamedAttrMap + }) as TensorInfo; + + const resultReshapeInputs: ReshapeInputs = {x: paddedXT}; + const resultReshapeAttrs: ReshapeAttrs = {shape: flattenShape}; + const result = reshapeConfig.kernelFunc({ + inputs: resultReshapeInputs, + backend, + attrs: resultReshapeAttrs as {} as NamedAttrMap + }); + + backend.disposeData(paddedX.dataId); + backend.disposeData(paddedXReshaped.dataId); + backend.disposeData(paddedXT.dataId); + + return result as TensorInfo; +} - return result; - } +export const spaceToBatchNDConfig: KernelConfig = { + kernelName: SpaceToBatchND, + backendName: 'cpu', + kernelFunc: spaceToBatchND as {} as KernelFunc }; diff --git a/tfjs-backend-cpu/src/kernels/Transpose_impl.ts b/tfjs-backend-cpu/src/kernels/Transpose_impl.ts index ffa2afa888a..3f7fbd2ddc4 100644 --- a/tfjs-backend-cpu/src/kernels/Transpose_impl.ts +++ b/tfjs-backend-cpu/src/kernels/Transpose_impl.ts @@ -30,7 +30,7 @@ export function transposeImpl( dtype as NumericDataType, util.sizeFromShape(newShape)); for (let i = 0; i < xSize; ++i) { - const loc = util.indexToLoc(i, xRank, xStrides); + const loc = util.indexToCoord(i, xRank, xStrides); // Permute location. const newLoc: number[] = new Array(loc.length); @@ -38,7 +38,7 @@ export function transposeImpl( newLoc[i] = loc[perm[i]]; } - const newIndex = util.locToIndex(newLoc, xRank, newStrides); + const newIndex = util.coordToIndex(newLoc, xRank, newStrides); result[newIndex] = xVals[i]; } return result; diff --git a/tfjs-backend-cpu/src/utils/kernel_utils.ts b/tfjs-backend-cpu/src/utils/kernel_utils.ts index 8c263201c15..cde975e121e 100644 --- a/tfjs-backend-cpu/src/utils/kernel_utils.ts +++ b/tfjs-backend-cpu/src/utils/kernel_utils.ts @@ -76,15 +76,15 @@ export function createBinaryKernelImpl(op: (a: number, b: number) => number) { } } else { for (let i = 0; i < result.length; ++i) { - const loc = util.indexToLoc(i, resultRank, resultStrides); + const loc = util.indexToCoord(i, resultRank, resultStrides); const aLoc = loc.slice(-aRank); aBroadcastDims.forEach(d => aLoc[d] = 0); - const aIndex = util.locToIndex(aLoc, aRank, aStrides); + const aIndex = util.coordToIndex(aLoc, aRank, aStrides); const bLoc = loc.slice(-bRank); bBroadcastDims.forEach(d => bLoc[d] = 0); - const bIndex = util.locToIndex(bLoc, bRank, bStrides); + const bIndex = util.coordToIndex(bLoc, bRank, bStrides); result[i] = op(aVals[aIndex], bVals[bIndex]); } diff --git a/tfjs-backend-wasm/src/backend_wasm.ts b/tfjs-backend-wasm/src/backend_wasm.ts index b1a7f4dc162..4b44fc7127d 100644 --- a/tfjs-backend-wasm/src/backend_wasm.ts +++ b/tfjs-backend-wasm/src/backend_wasm.ts @@ -59,6 +59,8 @@ export class BackendWasm extends KernelBackend { return this.dataIdMap.numDataIds(); } + decRef(dataId: DataId): void {} + async time(f: () => void): Promise { const start = util.now(); f(); diff --git a/tfjs-backend-webgl/src/backend_webgl.ts b/tfjs-backend-webgl/src/backend_webgl.ts index cf3c8d64258..35ebfb17bb4 100644 --- a/tfjs-backend-webgl/src/backend_webgl.ts +++ b/tfjs-backend-webgl/src/backend_webgl.ts @@ -267,6 +267,8 @@ export class MathBackendWebGL extends KernelBackend { this.pendingDeletes; } + decRef(dataId: DataId): void {} + write(values: BackendValues, shape: number[], dtype: DataType): DataId { if (env().getBool('WEBGL_CHECK_NUMERICAL_PROBLEMS') || env().getBool('DEBUG')) { diff --git a/tfjs-backend-webgpu/src/backend_webgpu.ts b/tfjs-backend-webgpu/src/backend_webgpu.ts index a27ab2ceb48..88784c87f68 100644 --- a/tfjs-backend-webgpu/src/backend_webgpu.ts +++ b/tfjs-backend-webgpu/src/backend_webgpu.ts @@ -1241,6 +1241,8 @@ export class WebGPUBackend extends KernelBackend { return this.tensorMap.numDataIds(); } + decRef(dataId: DataId): void {} + dispose() { if (this.disposed) { return; diff --git a/tfjs-core/src/backends/backend.ts b/tfjs-core/src/backends/backend.ts index 7fed43c2764..b5c11e51e3d 100644 --- a/tfjs-core/src/backends/backend.ts +++ b/tfjs-core/src/backends/backend.ts @@ -113,9 +113,6 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer { write(values: BackendValues, shape: number[], dtype: DataType): DataId { return notYetImplemented('write'); } - incRef(dataId: DataId): void { - return notYetImplemented('incRef'); - } decRef(dataId: DataId): void { return notYetImplemented('decRef'); } diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index 819b9db002f..0613b37738d 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -835,12 +835,10 @@ export class Engine implements TensorTracker, DataMover { } else { this.state.tensorInfo.get(a.dataId).refCount--; - // Todo(linazhao): Move decRef into backend's disposeData method once - // refCount mechanism is completely moved from Engine to backends. - // During the migration, non-modularized kernels will still use - // TensorInfo's refCounting whereas modularized kernels will use - // TensorData's refCounting, so we need to decRef both - // TensorInfo and TensorData's refCount to keep them in sync. + // Todo(linazhao): Remove decRef once reshape, clone and cast kernels + // are modularized. + // We added this to keep backend refCount to be in sync with engine + // refCount. info.backend.decRef(a.dataId); } // TODO(nsthorat): Construct an error and save the stack trace for diff --git a/tfjs-core/src/util.ts b/tfjs-core/src/util.ts index b5ead6f9150..737758f4117 100644 --- a/tfjs-core/src/util.ts +++ b/tfjs-core/src/util.ts @@ -732,14 +732,14 @@ export function decodeString(bytes: Uint8Array, encoding = 'utf-8'): string { } /** - * Computes flat index for a given location (multidimentionsal index) in a - * Tensor/multidimensional array. + * Computes flat index for a given coordinates in a Tensor/multidimensional + * array. * * @param locs Location in the tensor. * @param rank Rank of the tensor. * @param strides Tensor strides. */ -export function locToIndex( +export function coordToIndex( locs: number[], rank: number, strides: number[]): number { if (rank === 0) { return 0; @@ -754,14 +754,14 @@ export function locToIndex( } /** - * Computes the location (multidimensional index) in a tensor/multidimentional - * array for a given flat index. + * Computes the coordinates in a tensor/multidimentional array for a given flat + * index. * * @param index Index in flat array. * @param rank Rank of tensor. * @param strides Strides of tensor. */ -export function indexToLoc( +export function indexToCoord( index: number, rank: number, strides: number[]): number[] { if (rank === 0) { return []; diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index 466f68e30b8..41d5ec0802c 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -189,6 +189,8 @@ export class NodeJSKernelBackend extends KernelBackend { return this.tensorMap.numDataIds(); } + decRef(dataId: DataId): void {} + dispose(): void {} async read(dataId: DataId): Promise { From 56c986f2b71e22ff94c16a6cecb37549f17bf1b9 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 30 Jul 2020 20:45:04 -0700 Subject: [PATCH 08/22] Fix test. --- tfjs-core/src/engine_test.ts | 3 --- tfjs-core/src/kernel_registry_test.ts | 3 --- 2 files changed, 6 deletions(-) diff --git a/tfjs-core/src/engine_test.ts b/tfjs-core/src/engine_test.ts index 882a63b2c22..7bdcdc907a5 100644 --- a/tfjs-core/src/engine_test.ts +++ b/tfjs-core/src/engine_test.ts @@ -687,7 +687,6 @@ describeWithFlags('Detects memory leaks in kernels', ALL_ENVS, () => { dispose: () => null, disposeData: (dataId: {}) => null, numDataIds: () => dataIdsCount, - incRef: (dataId: {}) => null, decRef: (dataId: {}) => null } as TestStorage; }); @@ -715,7 +714,6 @@ describeWithFlags('Detects memory leaks in kernels', ALL_ENVS, () => { dispose: () => null, disposeData: (dataId: {}) => null, numDataIds: () => dataIdsCount, - incRef: (dataId: {}) => null, decRef: (dataId: {}) => null } as TestStorage; }); @@ -771,7 +769,6 @@ describe('Memory allocation outside a test scope', () => { read: async (dataId: object) => storedValues, dispose: () => null, disposeData: (dataId: {}) => null, - incRef: (dataId: {}) => {}, decRef: (dataId: {}) => {} } as TestStorage; }); diff --git a/tfjs-core/src/kernel_registry_test.ts b/tfjs-core/src/kernel_registry_test.ts index f799a2e5ee3..51d232524c2 100644 --- a/tfjs-core/src/kernel_registry_test.ts +++ b/tfjs-core/src/kernel_registry_test.ts @@ -101,7 +101,6 @@ describeWithFlags('kernel_registry', ALL_ENVS, () => { dispose: () => null, disposeData: (dataId: {}) => null, numDataIds: () => 0, - incRef: (dataId: {}) => null, decRef: (dataId: {}) => null } as TestBackend; }); @@ -111,7 +110,6 @@ describeWithFlags('kernel_registry', ALL_ENVS, () => { dispose: () => null, disposeData: (dataId: {}) => null, numDataIds: () => 0, - incRef: (dataId: {}) => null, decRef: (dataId: {}) => null } as TestBackend; }); @@ -153,7 +151,6 @@ describeWithFlags('kernel_registry', ALL_ENVS, () => { dispose: () => null, disposeData: (dataId: {}) => null, numDataIds: () => 0, - incRef: (dataId: {}) => null, decRef: (dataId: {}) => null } as TestBackend; tf.registerBackend(backendName, () => customBackend); From 892c4886e0549c2dc605784ae4ed03e431c8ba3e Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 30 Jul 2020 21:22:36 -0700 Subject: [PATCH 09/22] Fix test. --- tfjs-backend-cpu/src/backend_cpu.ts | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index a160a95ffb2..87bfb4409d9 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -100,6 +100,16 @@ export class MathBackendCPU extends KernelBackend { return dataId; } + /** Increase refCount of a `TensorData`. */ + incRef(dataId: DataId): void { + if (this.data.has(dataId)) { + const tensorData = this.data.get(dataId); + tensorData.refCount++; + } + // Do not throw error when dataId not found for testing. Some tests + // may use the backend without actually write any data to the backend. + } + /** Decrease refCount of a `TensorData`. */ decRef(dataId: DataId): void { if (this.data.has(dataId)) { From 5917635d1ca44477946bce20d868ee12b0a4bc99 Mon Sep 17 00:00:00 2001 From: Na Li Date: Fri, 31 Jul 2020 10:35:02 -0700 Subject: [PATCH 10/22] . --- tfjs-core/src/engine.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index 0613b37738d..518255fb133 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -839,7 +839,9 @@ export class Engine implements TensorTracker, DataMover { // are modularized. // We added this to keep backend refCount to be in sync with engine // refCount. - info.backend.decRef(a.dataId); + if (this.backendName === 'cpu') { + info.backend.decRef(a.dataId); + } } // TODO(nsthorat): Construct an error and save the stack trace for // debugging when in debug mode. Creating a stack trace is too expensive From 6af62446767805a5fd6b24e63e5989ce20093b68 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 3 Aug 2020 17:04:06 -0700 Subject: [PATCH 11/22] . --- tfjs-backend-cpu/src/backend_cpu.ts | 4 +- tfjs-backend-cpu/src/kernels/PadV2.ts | 4 +- tfjs-backend-cpu/src/kernels/Reshape.ts | 3 +- tfjs-backend-cpu/src/kernels/Reshape_test.ts | 33 +++++++++++++ .../src/kernels/SpaceToBatchND_test.ts | 48 +++++++++++++++++++ .../src/kernels/Transpose_impl.ts | 4 +- tfjs-backend-cpu/src/utils/kernel_utils.ts | 6 +-- tfjs-backend-wasm/src/backend_wasm.ts | 2 - tfjs-backend-webgl/src/backend_webgl.ts | 2 - tfjs-backend-webgpu/src/backend_webgpu.ts | 2 - tfjs-core/src/backends/backend.ts | 3 -- tfjs-core/src/engine.ts | 3 +- tfjs-core/src/engine_test.ts | 9 ++-- tfjs-core/src/kernel_registry_test.ts | 9 ++-- tfjs-core/src/ops/space_to_batch_nd_test.ts | 13 ----- tfjs-core/src/util.ts | 12 ++--- tfjs-node/src/nodejs_kernel_backend.ts | 2 - 17 files changed, 105 insertions(+), 54 deletions(-) create mode 100644 tfjs-backend-cpu/src/kernels/SpaceToBatchND_test.ts diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index 87bfb4409d9..1fa80d1d1de 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -105,9 +105,9 @@ export class MathBackendCPU extends KernelBackend { if (this.data.has(dataId)) { const tensorData = this.data.get(dataId); tensorData.refCount++; + } else { + throw new Error('incRef in cpu backend, but dataId not found.'); } - // Do not throw error when dataId not found for testing. Some tests - // may use the backend without actually write any data to the backend. } /** Decrease refCount of a `TensorData`. */ diff --git a/tfjs-backend-cpu/src/kernels/PadV2.ts b/tfjs-backend-cpu/src/kernels/PadV2.ts index df278551e37..731d0b381bd 100644 --- a/tfjs-backend-cpu/src/kernels/PadV2.ts +++ b/tfjs-backend-cpu/src/kernels/PadV2.ts @@ -50,9 +50,9 @@ function padV2( } for (let i = 0; i < xSize; i++) { - const coords = util.indexToCoord(i, xRank, xStrides); + const coords = util.indexToLoc(i, xRank, xStrides); const outCoords = coords.map((c, i) => c + start[i]); - const outIndex = util.coordToIndex(outCoords, resultRank, resultStrides); + const outIndex = util.locToIndex(outCoords, resultRank, resultStrides); resVals[outIndex] = xVals[i]; } diff --git a/tfjs-backend-cpu/src/kernels/Reshape.ts b/tfjs-backend-cpu/src/kernels/Reshape.ts index 968af215399..bc24b3384dc 100644 --- a/tfjs-backend-cpu/src/kernels/Reshape.ts +++ b/tfjs-backend-cpu/src/kernels/Reshape.ts @@ -25,9 +25,8 @@ function reshape( const {inputs, backend, attrs} = args; const {x} = inputs; const {shape} = attrs; - const cpuBackend = backend; - cpuBackend.incRef(x.dataId); + backend.incRef(x.dataId); return {dataId: x.dataId, shape, dtype: x.dtype}; } diff --git a/tfjs-backend-cpu/src/kernels/Reshape_test.ts b/tfjs-backend-cpu/src/kernels/Reshape_test.ts index 0d7903adea2..e92fb728819 100644 --- a/tfjs-backend-cpu/src/kernels/Reshape_test.ts +++ b/tfjs-backend-cpu/src/kernels/Reshape_test.ts @@ -42,4 +42,37 @@ describeWithFlags('Reshape.', ALL_ENVS, () => { const afterDisposeDataIds = tf.engine().backend.numDataIds(); expect(afterDisposeDataIds).toEqual(beforeDataIds); }); + + it('does not have memory leak calling reshape twice.', async () => { + const beforeDataIds = tf.engine().backend.numDataIds(); + + // Adding 1 new dataId. + const x = tf.tensor1d([1, 1, 1, 1]); + + // Does not add new dataId; + const res = + tf.engine().runKernel('Reshape', {x}, {shape: [2, 2]}) as Tensor; + + expectArraysEqual(res.shape, [2, 2]); + + // Does not add new dataId. + const res2 = + tf.engine().runKernel('Reshape', {x: res}, {shape: [1, 4]}) as Tensor; + expectArraysEqual(res2.shape, [1, 4]); + + const afterRes2DataIds = tf.engine().backend.numDataIds(); + expect(afterRes2DataIds).toEqual(beforeDataIds + 1); + + res.dispose(); + + const afterResDataIds = tf.engine().backend.numDataIds(); + expect(afterResDataIds).toEqual(beforeDataIds + 1); + + x.dispose(); + res2.dispose(); + + const afterDisposeDataIds = tf.engine().backend.numDataIds(); + // Should be able to dispose the dataId. + expect(afterDisposeDataIds).toEqual(beforeDataIds); + }); }); diff --git a/tfjs-backend-cpu/src/kernels/SpaceToBatchND_test.ts b/tfjs-backend-cpu/src/kernels/SpaceToBatchND_test.ts new file mode 100644 index 00000000000..f7766895f0d --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/SpaceToBatchND_test.ts @@ -0,0 +1,48 @@ +/** + * @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 * as tf from '@tensorflow/tfjs-core'; +import {test_util} from '@tensorflow/tfjs-core'; + +const {expectArraysClose} = test_util; +// tslint:disable-next-line: no-imports-from-dist +import {describeWithFlags, ALL_ENVS} from '@tensorflow/tfjs-core/dist/jasmine_util'; + +describeWithFlags('SpaceToBatchND.', ALL_ENVS, () => { + it('has no memory leak.', async () => { + const initialDataIds = tf.engine().backend.numDataIds(); + + const t = tf.tensor4d([[[[1], [2]], [[3], [4]]]], [1, 2, 2, 1]); + const blockShape = [2, 2]; + const paddings = [[0, 0], [0, 0]]; + + const res = tf.spaceToBatchND(t, blockShape, paddings); + expect(res.shape).toEqual([4, 1, 1, 1]); + expectArraysClose(await res.data(), [1, 2, 3, 4]); + + const afterResDataIds = tf.engine().backend.numDataIds(); + // 1 input tensor and 1 result tensor + expect(afterResDataIds).toEqual(initialDataIds + 2); + + t.dispose(); + res.dispose(); + + const afterDisposeDataIds = tf.engine().backend.numDataIds(); + + expect(afterDisposeDataIds).toEqual(initialDataIds); + }); +}); diff --git a/tfjs-backend-cpu/src/kernels/Transpose_impl.ts b/tfjs-backend-cpu/src/kernels/Transpose_impl.ts index 3f7fbd2ddc4..ffa2afa888a 100644 --- a/tfjs-backend-cpu/src/kernels/Transpose_impl.ts +++ b/tfjs-backend-cpu/src/kernels/Transpose_impl.ts @@ -30,7 +30,7 @@ export function transposeImpl( dtype as NumericDataType, util.sizeFromShape(newShape)); for (let i = 0; i < xSize; ++i) { - const loc = util.indexToCoord(i, xRank, xStrides); + const loc = util.indexToLoc(i, xRank, xStrides); // Permute location. const newLoc: number[] = new Array(loc.length); @@ -38,7 +38,7 @@ export function transposeImpl( newLoc[i] = loc[perm[i]]; } - const newIndex = util.coordToIndex(newLoc, xRank, newStrides); + const newIndex = util.locToIndex(newLoc, xRank, newStrides); result[newIndex] = xVals[i]; } return result; diff --git a/tfjs-backend-cpu/src/utils/kernel_utils.ts b/tfjs-backend-cpu/src/utils/kernel_utils.ts index cde975e121e..8c263201c15 100644 --- a/tfjs-backend-cpu/src/utils/kernel_utils.ts +++ b/tfjs-backend-cpu/src/utils/kernel_utils.ts @@ -76,15 +76,15 @@ export function createBinaryKernelImpl(op: (a: number, b: number) => number) { } } else { for (let i = 0; i < result.length; ++i) { - const loc = util.indexToCoord(i, resultRank, resultStrides); + const loc = util.indexToLoc(i, resultRank, resultStrides); const aLoc = loc.slice(-aRank); aBroadcastDims.forEach(d => aLoc[d] = 0); - const aIndex = util.coordToIndex(aLoc, aRank, aStrides); + const aIndex = util.locToIndex(aLoc, aRank, aStrides); const bLoc = loc.slice(-bRank); bBroadcastDims.forEach(d => bLoc[d] = 0); - const bIndex = util.coordToIndex(bLoc, bRank, bStrides); + const bIndex = util.locToIndex(bLoc, bRank, bStrides); result[i] = op(aVals[aIndex], bVals[bIndex]); } diff --git a/tfjs-backend-wasm/src/backend_wasm.ts b/tfjs-backend-wasm/src/backend_wasm.ts index 4b44fc7127d..b1a7f4dc162 100644 --- a/tfjs-backend-wasm/src/backend_wasm.ts +++ b/tfjs-backend-wasm/src/backend_wasm.ts @@ -59,8 +59,6 @@ export class BackendWasm extends KernelBackend { return this.dataIdMap.numDataIds(); } - decRef(dataId: DataId): void {} - async time(f: () => void): Promise { const start = util.now(); f(); diff --git a/tfjs-backend-webgl/src/backend_webgl.ts b/tfjs-backend-webgl/src/backend_webgl.ts index 35ebfb17bb4..cf3c8d64258 100644 --- a/tfjs-backend-webgl/src/backend_webgl.ts +++ b/tfjs-backend-webgl/src/backend_webgl.ts @@ -267,8 +267,6 @@ export class MathBackendWebGL extends KernelBackend { this.pendingDeletes; } - decRef(dataId: DataId): void {} - write(values: BackendValues, shape: number[], dtype: DataType): DataId { if (env().getBool('WEBGL_CHECK_NUMERICAL_PROBLEMS') || env().getBool('DEBUG')) { diff --git a/tfjs-backend-webgpu/src/backend_webgpu.ts b/tfjs-backend-webgpu/src/backend_webgpu.ts index 88784c87f68..a27ab2ceb48 100644 --- a/tfjs-backend-webgpu/src/backend_webgpu.ts +++ b/tfjs-backend-webgpu/src/backend_webgpu.ts @@ -1241,8 +1241,6 @@ export class WebGPUBackend extends KernelBackend { return this.tensorMap.numDataIds(); } - decRef(dataId: DataId): void {} - dispose() { if (this.disposed) { return; diff --git a/tfjs-core/src/backends/backend.ts b/tfjs-core/src/backends/backend.ts index b5c11e51e3d..66907f2fa8d 100644 --- a/tfjs-core/src/backends/backend.ts +++ b/tfjs-core/src/backends/backend.ts @@ -113,9 +113,6 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer { write(values: BackendValues, shape: number[], dtype: DataType): DataId { return notYetImplemented('write'); } - decRef(dataId: DataId): void { - return notYetImplemented('decRef'); - } move(dataId: DataId, values: BackendValues, shape: number[], dtype: DataType): void { return notYetImplemented('move'); diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index 518255fb133..c0ea6bfe6c5 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -840,7 +840,8 @@ export class Engine implements TensorTracker, DataMover { // We added this to keep backend refCount to be in sync with engine // refCount. if (this.backendName === 'cpu') { - info.backend.decRef(a.dataId); + // tslint:disable-next-line: no-any + (info.backend as any).decRef(a.dataId); } } // TODO(nsthorat): Construct an error and save the stack trace for diff --git a/tfjs-core/src/engine_test.ts b/tfjs-core/src/engine_test.ts index 7bdcdc907a5..f6178b2103f 100644 --- a/tfjs-core/src/engine_test.ts +++ b/tfjs-core/src/engine_test.ts @@ -686,8 +686,7 @@ describeWithFlags('Detects memory leaks in kernels', ALL_ENVS, () => { id: 1, dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => dataIdsCount, - decRef: (dataId: {}) => null + numDataIds: () => dataIdsCount } as TestStorage; }); @@ -713,8 +712,7 @@ describeWithFlags('Detects memory leaks in kernels', ALL_ENVS, () => { id: 1, dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => dataIdsCount, - decRef: (dataId: {}) => null + numDataIds: () => dataIdsCount } as TestStorage; }); tf.setBackend(backendName); @@ -768,8 +766,7 @@ describe('Memory allocation outside a test scope', () => { }, read: async (dataId: object) => storedValues, dispose: () => null, - disposeData: (dataId: {}) => null, - decRef: (dataId: {}) => {} + disposeData: (dataId: {}) => null } as TestStorage; }); tf.setBackend(backendName); diff --git a/tfjs-core/src/kernel_registry_test.ts b/tfjs-core/src/kernel_registry_test.ts index 51d232524c2..352fce14391 100644 --- a/tfjs-core/src/kernel_registry_test.ts +++ b/tfjs-core/src/kernel_registry_test.ts @@ -100,8 +100,7 @@ describeWithFlags('kernel_registry', ALL_ENVS, () => { id: 1, dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => 0, - decRef: (dataId: {}) => null + numDataIds: () => 0 } as TestBackend; }); tf.registerBackend('backend2', () => { @@ -109,8 +108,7 @@ describeWithFlags('kernel_registry', ALL_ENVS, () => { id: 2, dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => 0, - decRef: (dataId: {}) => null + numDataIds: () => 0 } as TestBackend; }); @@ -150,8 +148,7 @@ describeWithFlags('kernel_registry', ALL_ENVS, () => { const customBackend = { dispose: () => null, disposeData: (dataId: {}) => null, - numDataIds: () => 0, - decRef: (dataId: {}) => null + numDataIds: () => 0 } as TestBackend; tf.registerBackend(backendName, () => customBackend); diff --git a/tfjs-core/src/ops/space_to_batch_nd_test.ts b/tfjs-core/src/ops/space_to_batch_nd_test.ts index f6485ec8f25..11da95ac170 100644 --- a/tfjs-core/src/ops/space_to_batch_nd_test.ts +++ b/tfjs-core/src/ops/space_to_batch_nd_test.ts @@ -21,8 +21,6 @@ import {expectArraysClose} from '../test_util'; describeWithFlags('spaceToBatchND', ALL_ENVS, () => { it('tensor4d, input shape=[1, 2, 2, 1], blockShape=[2, 2]', async () => { - const initialDataIds = tf.engine().backend.numDataIds(); - const t = tf.tensor4d([[[[1], [2]], [[3], [4]]]], [1, 2, 2, 1]); const blockShape = [2, 2]; const paddings = [[0, 0], [0, 0]]; @@ -30,17 +28,6 @@ describeWithFlags('spaceToBatchND', ALL_ENVS, () => { const res = tf.spaceToBatchND(t, blockShape, paddings); expect(res.shape).toEqual([4, 1, 1, 1]); expectArraysClose(await res.data(), [1, 2, 3, 4]); - - const afterResDataIds = tf.engine().backend.numDataIds(); - // 1 input tensor and 1 result tensor - expect(afterResDataIds).toEqual(initialDataIds + 2); - - t.dispose(); - res.dispose(); - - const afterDisposeDataIds = tf.engine().backend.numDataIds(); - - expect(afterDisposeDataIds).toEqual(initialDataIds); }); it('tensor4d, input shape=[1, 2, 2, 3], blockShape=[2, 2]', async () => { diff --git a/tfjs-core/src/util.ts b/tfjs-core/src/util.ts index 737758f4117..b5ead6f9150 100644 --- a/tfjs-core/src/util.ts +++ b/tfjs-core/src/util.ts @@ -732,14 +732,14 @@ export function decodeString(bytes: Uint8Array, encoding = 'utf-8'): string { } /** - * Computes flat index for a given coordinates in a Tensor/multidimensional - * array. + * Computes flat index for a given location (multidimentionsal index) in a + * Tensor/multidimensional array. * * @param locs Location in the tensor. * @param rank Rank of the tensor. * @param strides Tensor strides. */ -export function coordToIndex( +export function locToIndex( locs: number[], rank: number, strides: number[]): number { if (rank === 0) { return 0; @@ -754,14 +754,14 @@ export function coordToIndex( } /** - * Computes the coordinates in a tensor/multidimentional array for a given flat - * index. + * Computes the location (multidimensional index) in a tensor/multidimentional + * array for a given flat index. * * @param index Index in flat array. * @param rank Rank of tensor. * @param strides Strides of tensor. */ -export function indexToCoord( +export function indexToLoc( index: number, rank: number, strides: number[]): number[] { if (rank === 0) { return []; diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index 41d5ec0802c..466f68e30b8 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -189,8 +189,6 @@ export class NodeJSKernelBackend extends KernelBackend { return this.tensorMap.numDataIds(); } - decRef(dataId: DataId): void {} - dispose(): void {} async read(dataId: DataId): Promise { From 6bd3eb4d2145d193d8b02dc0f35d0b8ade100691 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 3 Aug 2020 17:05:27 -0700 Subject: [PATCH 12/22] . --- tfjs-backend-cpu/src/backend_cpu.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index 1fa80d1d1de..bff434c2c67 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -116,8 +116,6 @@ export class MathBackendCPU extends KernelBackend { const tensorData = this.data.get(dataId); tensorData.refCount--; } - // Do not throw error when dataId not found for testing. Some tests - // may use the backend without actually write any data to the backend. } move( From 61955555614568ffafaa1e17555c22d88ae9dc40 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 3 Aug 2020 17:07:23 -0700 Subject: [PATCH 13/22] . --- tfjs-backend-cpu/src/backend_cpu_test.ts | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/tfjs-backend-cpu/src/backend_cpu_test.ts b/tfjs-backend-cpu/src/backend_cpu_test.ts index 25ed315ee6a..2c78b6dc8bb 100644 --- a/tfjs-backend-cpu/src/backend_cpu_test.ts +++ b/tfjs-backend-cpu/src/backend_cpu_test.ts @@ -132,24 +132,4 @@ describeWithFlags('memory cpu', ALL_ENVS, () => { '(2 bytes per character)'; expect(mem.reasons.indexOf(expectedReasonString) >= 0).toBe(true); }); - - it('does not have memory leak with reshape.', async () => { - const beforeDataIds = tf.engine().backend.numDataIds(); - - const x = tf.tensor1d([1, 1, 1, 1]); - const res = - tf.engine().runKernel('Reshape', {x}, {shape: [2, 2]}) as Tensor; - - expectArraysClose(await res.data(), [1, 1, 1, 1]); - expectArraysEqual(res.shape, [2, 2]); - - const afterResDataIds = tf.engine().backend.numDataIds(); - expect(afterResDataIds).toEqual(beforeDataIds + 1); - - x.dispose(); - res.dispose(); - - const afterDisposeDataIds = tf.engine().backend.numDataIds(); - expect(afterDisposeDataIds).toEqual(beforeDataIds); - }); }); From 13b2b60d5965afc877d36022373efd6bf18f8600 Mon Sep 17 00:00:00 2001 From: Na Li Date: Mon, 3 Aug 2020 17:22:52 -0700 Subject: [PATCH 14/22] . --- tfjs-backend-cpu/src/backend_cpu_test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tfjs-backend-cpu/src/backend_cpu_test.ts b/tfjs-backend-cpu/src/backend_cpu_test.ts index 2c78b6dc8bb..b4aa455dfdc 100644 --- a/tfjs-backend-cpu/src/backend_cpu_test.ts +++ b/tfjs-backend-cpu/src/backend_cpu_test.ts @@ -16,9 +16,9 @@ */ import * as tf from '@tensorflow/tfjs-core'; -import {engine, Tensor, test_util, util} from '@tensorflow/tfjs-core'; +import {engine, test_util, util} from '@tensorflow/tfjs-core'; -const {expectArraysClose, expectArraysEqual} = test_util; +const {expectArraysEqual} = test_util; // tslint:disable-next-line: no-imports-from-dist import {describeWithFlags, ALL_ENVS} from '@tensorflow/tfjs-core/dist/jasmine_util'; From 1a8cbc1a70716872b38c066c4a936bf16db63e63 Mon Sep 17 00:00:00 2001 From: Na Li Date: Tue, 4 Aug 2020 23:05:40 -0700 Subject: [PATCH 15/22] Add move data logic. --- e2e/integration_tests/backends_test.ts | 103 ++++++++++++++++++++++++ tfjs-backend-cpu/src/backend_cpu.ts | 11 ++- tfjs-backend-cpu/src/kernels/Reshape.ts | 3 + tfjs-backend-webgl/src/backend_webgl.ts | 1 + tfjs-core/src/backends/backend.ts | 2 +- tfjs-core/src/engine.ts | 3 +- 6 files changed, 119 insertions(+), 4 deletions(-) create mode 100644 e2e/integration_tests/backends_test.ts diff --git a/e2e/integration_tests/backends_test.ts b/e2e/integration_tests/backends_test.ts new file mode 100644 index 00000000000..ac0f08dc27e --- /dev/null +++ b/e2e/integration_tests/backends_test.ts @@ -0,0 +1,103 @@ +/** + * @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 '@tensorflow/tfjs-backend-cpu'; +import '@tensorflow/tfjs-backend-webgl'; + +import * as tfc from '@tensorflow/tfjs-core'; + +import {SMOKE} from './constants'; + +/** + * This file tests backend switching scenario. + */ +describe(`${SMOKE} backends`, () => { + describe('switch', () => { + beforeAll(() => { + tfc.env().set('WEBGL_CPU_FORWARD', false); + }); + + it(`from webgl to cpu.`, async () => { + // A backend with no refCounter. + await tfc.setBackend('webgl'); + + // numDataIds = 0 + const before = tfc.engine().backend.numDataIds(); + + // input tensorData is stored in webgl backend. + const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32'); + + // input tensorInfo refCount = 2; + const inputReshaped = tfc.reshape(input, [2, 2]); + + await tfc.setBackend('cpu'); + + // inputReshaped tensorData is reshaped, during which it is moved to cpu + // backend with a refCount 1. After reshape, it has a refCount 2. The + // tensorInfo refCount becomes 3 because inputReshaped2 is an output. + const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]); + + input.dispose(); + inputReshaped.dispose(); + inputReshaped2.dispose(); + + const after = tfc.engine().backend.numDataIds(); + + expect(after).toBe(before); + + // There should be no more data in webgl, because they are all moved to + // cpu even if the tensorInfo refCount before moving is 2. + await tfc.setBackend('webgl'); + const webglAfter = tfc.engine().backend.numDataIds(); + expect(webglAfter).toBe(0); + }); + + it(`from cpu to webgl.`, async () => { + await tfc.setBackend('cpu'); + + // numDataIds = 0. + const before = tfc.engine().backend.numDataIds(); + + // input tensorData is stored in cpu backend. + const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32'); + + // input tensorData refCount = 2, tensorInfo refCount = 1. + const inputReshaped = tfc.reshape(input, [2, 2]); + + await tfc.setBackend('webgl'); + + // inputReshaped tensorData is reshaped, during which it is moved to webgl + // backend with a refCount 2. Then tensorInfo refCount is also 2. After + // reshape, the tensorInfo refCount becomes 3. + const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]); + + input.dispose(); + inputReshaped.dispose(); + inputReshaped2.dispose(); + + const after = tfc.engine().backend.numDataIds(); + + expect(after).toBe(before); + + // There should be no more data in cpu, because they are all moved to + // webgl. + await tfc.setBackend('cpu'); + const cpuAfter = tfc.engine().backend.numDataIds(); + expect(cpuAfter).toBe(0); + }); + }); +}); diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index bff434c2c67..916594bd2f7 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -163,13 +163,15 @@ export class MathBackendCPU extends KernelBackend { return engine().makeTensorFromDataId(dataId, shape, dtype, this) as T; } - disposeData(dataId: DataId): void { + disposeData(dataId: DataId, force?: boolean): void { if (this.data.has(dataId)) { const tensorData = this.data.get(dataId); tensorData.refCount--; - if (tensorData.refCount < 1) { + const shouldDelete = force || (tensorData.refCount < 1); + + if (shouldDelete) { if (tensorData.complexTensors != null) { // Todo(linazhao): Change to disposeData once complex, real, and imag // kernels are modularized and real and imag becomes `TensorInfo`. @@ -435,6 +437,7 @@ export class MathBackendCPU extends KernelBackend { batchMatMul( a: Tensor3D, b: Tensor3D, transposeA: boolean, transposeB: boolean): Tensor3D { + console.log('IN CPU BATCHMATMUL'); assertNotComplex([a, b], 'matMul'); const sharedDim = transposeA ? a.shape[1] : a.shape[2]; @@ -486,15 +489,19 @@ export class MathBackendCPU extends KernelBackend { fusedBatchMatMul( {a, b, transposeA, transposeB, bias, activation, preluActivationWeights}: backend_util.FusedBatchMatMulConfig): Tensor3D { + console.log('IM IN FUSED BATCH MAT MUL'); let result = this.batchMatMul(a, b, transposeA, transposeB); if (bias) { + console.log('GOING TO ADD'); result = this.add(result, bias) as Tensor3D; } if (activation) { + console.log('GOING TO map activation'); result = mapActivation(this, result, activation, preluActivationWeights) as Tensor3D; } + console.log('RETURN RESULT'); return result; } diff --git a/tfjs-backend-cpu/src/kernels/Reshape.ts b/tfjs-backend-cpu/src/kernels/Reshape.ts index bc24b3384dc..af9f83db131 100644 --- a/tfjs-backend-cpu/src/kernels/Reshape.ts +++ b/tfjs-backend-cpu/src/kernels/Reshape.ts @@ -26,6 +26,9 @@ function reshape( const {x} = inputs; const {shape} = attrs; + // Make sure data is in cpu backend. + backend.data.get(x.dataId); + backend.incRef(x.dataId); return {dataId: x.dataId, shape, dtype: x.dtype}; diff --git a/tfjs-backend-webgl/src/backend_webgl.ts b/tfjs-backend-webgl/src/backend_webgl.ts index cf3c8d64258..7c67b71ea93 100644 --- a/tfjs-backend-webgl/src/backend_webgl.ts +++ b/tfjs-backend-webgl/src/backend_webgl.ts @@ -2200,6 +2200,7 @@ export class MathBackendWebGL extends KernelBackend { reshape(x: Tensor, shape: ShapeMap[R]): Tensor { const texData = this.texData.get(x.dataId); + console.log('IN WEBGL RESHAPE'); if (texData.isPacked && !webgl_util.isReshapeFree(x.shape, shape) && !(texData.texture !== null && webgl_util.isReshapeFree(texData.shape, shape))) { diff --git a/tfjs-core/src/backends/backend.ts b/tfjs-core/src/backends/backend.ts index 66907f2fa8d..47aec0261b9 100644 --- a/tfjs-core/src/backends/backend.ts +++ b/tfjs-core/src/backends/backend.ts @@ -107,7 +107,7 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer { numDataIds(): number { return notYetImplemented('numDataIds'); } - disposeData(dataId: object): void { + disposeData(dataId: object, force?: boolean): void { return notYetImplemented('disposeData'); } write(values: BackendValues, shape: number[], dtype: DataType): DataId { diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index c0ea6bfe6c5..b5d9beaf7e5 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -387,7 +387,7 @@ export class Engine implements TensorTracker, DataMover { const values = this.readSync(dataId); // Delete the tensor from the old backend and move it to the new // backend. - srcBackend.disposeData(dataId); + srcBackend.disposeData(dataId, true /* force */); info.backend = backend; backend.move(dataId, values, info.shape, info.dtype); if (this.shouldCheckForMemLeaks()) { @@ -822,6 +822,7 @@ export class Engine implements TensorTracker, DataMover { const info = this.state.tensorInfo.get(a.dataId); const refCount = info.refCount; + console.log(`refCount is: ${refCount}`); if (refCount <= 1) { // Don't count bytes for complex numbers as they are counted by their // components. From 243f23ce5a558178f70242b12378752e1857fed7 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 6 Aug 2020 08:58:53 -0700 Subject: [PATCH 16/22] Add custom op test. --- e2e/integration_tests/custom_op_test.ts | 82 +++++++++++++++++++ tfjs-backend-cpu/src/backend_cpu.ts | 42 +++++----- tfjs-backend-cpu/src/kernels/Reshape.ts | 3 - .../src/kernels/SpaceToBatchND.ts | 6 +- tfjs-backend-webgl/src/backend_webgl.ts | 2 +- tfjs-core/src/backends/backend.ts | 2 +- tfjs-core/src/engine.ts | 3 +- 7 files changed, 108 insertions(+), 32 deletions(-) create mode 100644 e2e/integration_tests/custom_op_test.ts diff --git a/e2e/integration_tests/custom_op_test.ts b/e2e/integration_tests/custom_op_test.ts new file mode 100644 index 00000000000..aa25905b625 --- /dev/null +++ b/e2e/integration_tests/custom_op_test.ts @@ -0,0 +1,82 @@ +import '@tensorflow/tfjs-backend-cpu'; + +import * as tfconverter from '@tensorflow/tfjs-converter'; +import * as tfc from '@tensorflow/tfjs-core'; + +import {SMOKE} from './constants'; + +const HOST = 'http://example.org'; +const MODEL_URL = `${HOST}/model.json`; + +const CUSTOM_OP_MODEL = { + node: [ + { + name: 'Input', + op: 'Placeholder', + attr: { + dtype: { + type: 1, // DT_FLOAT + }, + shape: {shape: {dim: [{size: 4}]}} + } + }, + {name: 'CustomOp', op: 'CustomOp', input: ['Input'], attr: {}} + ], + versions: {producer: 1.0, minConsumer: 3} +}; + +const weightsManifest: tfc.io.WeightsManifestEntry[] = + [{'name': 'Const', 'dtype': 'float32', 'shape': [1]}]; + +const bias = tfc.tensor1d([0], 'float32'); + +const CUSTOM_HTTP_MODEL_LOADER = { + load: async () => { + return { + modelTopology: CUSTOM_OP_MODEL, + weightSpecs: weightsManifest, + weightData: bias.dataSync(), + format: 'tfjs-graph-model', + generatedBy: '1.15', + convertedBy: '1.3.1' + }; + } +}; + +describe(`${SMOKE} custom model`, () => { + it('refCounter works correctly.', async () => { + const model = new tfconverter.GraphModel(MODEL_URL); + + spyOn(tfc.io, 'getLoadHandlers').and.returnValue([ + CUSTOM_HTTP_MODEL_LOADER + ]); + + // A custom op that calls unmodularized kernels and modularized kernels. + // softmax ker + tfconverter.registerOp('CustomOp', (nodeValue) => { + const x = nodeValue.inputs[0]; + const softMax = tfc.softmax(x); + const clone = tfc.clone(softMax); + return [tfc.reshape(clone, [2, 2])]; + }); + + await model.load(); + + const before = tfc.memory().numTensors; + + const input = tfc.tensor1d([1, 2, 3, 4]); + const output = model.predict(input) as tfc.Tensor; + + tfc.test_util.expectArraysClose(await output.data(), [ + 0.032058604061603546, 0.08714432269334793, 0.23688283562660217, + 0.6439142823219299 + ]); + + input.dispose(); + output.dispose(); + + const after = tfc.memory().numTensors; + + expect(after).toEqual(before); + }); +}); diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index 916594bd2f7..0b15b45b037 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -102,12 +102,8 @@ export class MathBackendCPU extends KernelBackend { /** Increase refCount of a `TensorData`. */ incRef(dataId: DataId): void { - if (this.data.has(dataId)) { - const tensorData = this.data.get(dataId); - tensorData.refCount++; - } else { - throw new Error('incRef in cpu backend, but dataId not found.'); - } + const tensorData = this.data.get(dataId); + tensorData.refCount++; } /** Decrease refCount of a `TensorData`. */ @@ -163,23 +159,29 @@ export class MathBackendCPU extends KernelBackend { return engine().makeTensorFromDataId(dataId, shape, dtype, this) as T; } - disposeData(dataId: DataId, force?: boolean): void { + disposeData(dataId: DataId): void { if (this.data.has(dataId)) { const tensorData = this.data.get(dataId); - tensorData.refCount--; + if (tensorData.complexTensors != null) { + // Todo(linazhao): Change to disposeData once complex, real, and imag + // kernels are modularized and real and imag becomes `TensorInfo`. + tensorData.complexTensors.real.dispose(); + tensorData.complexTensors.imag.dispose(); + } + + this.data.delete(dataId); + } + } - const shouldDelete = force || (tensorData.refCount < 1); + disposeDataSoft(dataId: DataId): void { + if (this.data.has(dataId)) { + const tensorData = this.data.get(dataId); - if (shouldDelete) { - if (tensorData.complexTensors != null) { - // Todo(linazhao): Change to disposeData once complex, real, and imag - // kernels are modularized and real and imag becomes `TensorInfo`. - tensorData.complexTensors.real.dispose(); - tensorData.complexTensors.imag.dispose(); - } + tensorData.refCount--; - this.data.delete(dataId); + if (tensorData.refCount < 1) { + this.disposeData(dataId); } } } @@ -437,7 +439,6 @@ export class MathBackendCPU extends KernelBackend { batchMatMul( a: Tensor3D, b: Tensor3D, transposeA: boolean, transposeB: boolean): Tensor3D { - console.log('IN CPU BATCHMATMUL'); assertNotComplex([a, b], 'matMul'); const sharedDim = transposeA ? a.shape[1] : a.shape[2]; @@ -489,19 +490,16 @@ export class MathBackendCPU extends KernelBackend { fusedBatchMatMul( {a, b, transposeA, transposeB, bias, activation, preluActivationWeights}: backend_util.FusedBatchMatMulConfig): Tensor3D { - console.log('IM IN FUSED BATCH MAT MUL'); let result = this.batchMatMul(a, b, transposeA, transposeB); if (bias) { - console.log('GOING TO ADD'); result = this.add(result, bias) as Tensor3D; } if (activation) { - console.log('GOING TO map activation'); result = mapActivation(this, result, activation, preluActivationWeights) as Tensor3D; } - console.log('RETURN RESULT'); + return result; } diff --git a/tfjs-backend-cpu/src/kernels/Reshape.ts b/tfjs-backend-cpu/src/kernels/Reshape.ts index af9f83db131..bc24b3384dc 100644 --- a/tfjs-backend-cpu/src/kernels/Reshape.ts +++ b/tfjs-backend-cpu/src/kernels/Reshape.ts @@ -26,9 +26,6 @@ function reshape( const {x} = inputs; const {shape} = attrs; - // Make sure data is in cpu backend. - backend.data.get(x.dataId); - backend.incRef(x.dataId); return {dataId: x.dataId, shape, dtype: x.dtype}; diff --git a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts index 4fa4eb3e4b0..3553d6d0327 100644 --- a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts +++ b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts @@ -84,9 +84,9 @@ function spaceToBatchND(args: { attrs: resultReshapeAttrs as {} as NamedAttrMap }); - backend.disposeData(paddedX.dataId); - backend.disposeData(paddedXReshaped.dataId); - backend.disposeData(paddedXT.dataId); + backend.disposeDataSoft(paddedX.dataId); + backend.disposeDataSoft(paddedXReshaped.dataId); + backend.disposeDataSoft(paddedXT.dataId); return result as TensorInfo; } diff --git a/tfjs-backend-webgl/src/backend_webgl.ts b/tfjs-backend-webgl/src/backend_webgl.ts index 7c67b71ea93..a939a659304 100644 --- a/tfjs-backend-webgl/src/backend_webgl.ts +++ b/tfjs-backend-webgl/src/backend_webgl.ts @@ -2200,7 +2200,7 @@ export class MathBackendWebGL extends KernelBackend { reshape(x: Tensor, shape: ShapeMap[R]): Tensor { const texData = this.texData.get(x.dataId); - console.log('IN WEBGL RESHAPE'); + if (texData.isPacked && !webgl_util.isReshapeFree(x.shape, shape) && !(texData.texture !== null && webgl_util.isReshapeFree(texData.shape, shape))) { diff --git a/tfjs-core/src/backends/backend.ts b/tfjs-core/src/backends/backend.ts index 47aec0261b9..66907f2fa8d 100644 --- a/tfjs-core/src/backends/backend.ts +++ b/tfjs-core/src/backends/backend.ts @@ -107,7 +107,7 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer { numDataIds(): number { return notYetImplemented('numDataIds'); } - disposeData(dataId: object, force?: boolean): void { + disposeData(dataId: object): void { return notYetImplemented('disposeData'); } write(values: BackendValues, shape: number[], dtype: DataType): DataId { diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index b5d9beaf7e5..c0ea6bfe6c5 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -387,7 +387,7 @@ export class Engine implements TensorTracker, DataMover { const values = this.readSync(dataId); // Delete the tensor from the old backend and move it to the new // backend. - srcBackend.disposeData(dataId, true /* force */); + srcBackend.disposeData(dataId); info.backend = backend; backend.move(dataId, values, info.shape, info.dtype); if (this.shouldCheckForMemLeaks()) { @@ -822,7 +822,6 @@ export class Engine implements TensorTracker, DataMover { const info = this.state.tensorInfo.get(a.dataId); const refCount = info.refCount; - console.log(`refCount is: ${refCount}`); if (refCount <= 1) { // Don't count bytes for complex numbers as they are counted by their // components. From 28cbe2d8eb054174cdcda6b14ccdca2c02e72bfd Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 6 Aug 2020 12:41:30 -0700 Subject: [PATCH 17/22] Refactor test. --- e2e/integration_tests/backends_test.ts | 96 ++++++++++++++++++-------- 1 file changed, 69 insertions(+), 27 deletions(-) diff --git a/e2e/integration_tests/backends_test.ts b/e2e/integration_tests/backends_test.ts index ac0f08dc27e..f1a4704b2c4 100644 --- a/e2e/integration_tests/backends_test.ts +++ b/e2e/integration_tests/backends_test.ts @@ -35,69 +35,111 @@ describe(`${SMOKE} backends`, () => { // A backend with no refCounter. await tfc.setBackend('webgl'); - // numDataIds = 0 - const before = tfc.engine().backend.numDataIds(); + const webglBefore = tfc.engine().backend.numDataIds(); - // input tensorData is stored in webgl backend. + // input is stored in webgl backend. tensorInfo refCount = 1. const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32'); // input tensorInfo refCount = 2; const inputReshaped = tfc.reshape(input, [2, 2]); + const webglAfter = tfc.engine().backend.numDataIds(); + + expect(webglAfter).toEqual(webglBefore + 1); + await tfc.setBackend('cpu'); - // inputReshaped tensorData is reshaped, during which it is moved to cpu - // backend with a refCount 1. After reshape, it has a refCount 2. The - // tensorInfo refCount becomes 3 because inputReshaped2 is an output. + const cpuBefore = tfc.engine().backend.numDataIds(); + + // input moved to cpu. tensorData refCount = 1. const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]); + const cpuAfter = tfc.engine().backend.numDataIds(); + + // input tensorData refCount = 3, reshape increase by 1, then output + // tensor increase by 1. + expect(cpuAfter).toEqual(cpuBefore + 1); + + await tfc.setBackend('webgl'); + + // Because input is moved to cpu, data should be deleted from webgl. + expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter - 1); + + await tfc.setBackend('cpu'); + + // After dispose, tensorData refCount = 2. input.dispose(); + + // Input is not deleted, because refCount is 2. + expect(tfc.engine().backend.numDataIds()).toEqual(cpuAfter); + + // After dispose, tensorData refCount = 1. inputReshaped.dispose(); + + // Input is not deleted, because refCount is 1. + expect(tfc.engine().backend.numDataIds()).toEqual(cpuAfter); + + // After dispose, tensorData refCount = 0, data deleted. inputReshaped2.dispose(); const after = tfc.engine().backend.numDataIds(); - expect(after).toBe(before); - - // There should be no more data in webgl, because they are all moved to - // cpu even if the tensorInfo refCount before moving is 2. - await tfc.setBackend('webgl'); - const webglAfter = tfc.engine().backend.numDataIds(); - expect(webglAfter).toBe(0); + expect(after).toBe(cpuBefore); }); it(`from cpu to webgl.`, async () => { await tfc.setBackend('cpu'); - // numDataIds = 0. - const before = tfc.engine().backend.numDataIds(); + const cpuBefore = tfc.engine().backend.numDataIds(); - // input tensorData is stored in cpu backend. + // input is stored in cpu backend. tensorData refCount = 1, tensorInfo + // refCount = 1. const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32'); - // input tensorData refCount = 2, tensorInfo refCount = 1. + // input tensorData refCount = 3, reshape increase by 1, then output + // tensor increase by 1. tensorInfo refCount = 1. const inputReshaped = tfc.reshape(input, [2, 2]); + const cpuAfter = tfc.engine().backend.numDataIds(); + + expect(cpuAfter).toEqual(cpuBefore + 1); + await tfc.setBackend('webgl'); - // inputReshaped tensorData is reshaped, during which it is moved to webgl - // backend with a refCount 2. Then tensorInfo refCount is also 2. After - // reshape, the tensorInfo refCount becomes 3. + const webglBefore = tfc.engine().backend.numDataIds(); + + // input moved to webgl. tensorInfo refCount = 3. const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]); + const webglAfter = tfc.engine().backend.numDataIds(); + + expect(webglAfter).toEqual(webglBefore + 1); + + await tfc.setBackend('cpu'); + + // Because input is moved to webgl, data should be deleted from cpu. + expect(tfc.engine().backend.numDataIds()).toEqual(cpuAfter - 1); + + await tfc.setBackend('webgl'); + + // After dispose, tensorInfo = 2. input.dispose(); + + // Data is not deleted, because tensorInfo is 2. + expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter); + + // After dipose, tensorInfo = 1. inputReshaped.dispose(); + + // Data is not deleted, because tensorInfo is 1. + expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter); + + // After dipose, tensorInfo = 1, data deleted. inputReshaped2.dispose(); const after = tfc.engine().backend.numDataIds(); - expect(after).toBe(before); - - // There should be no more data in cpu, because they are all moved to - // webgl. - await tfc.setBackend('cpu'); - const cpuAfter = tfc.engine().backend.numDataIds(); - expect(cpuAfter).toBe(0); + expect(after).toBe(webglBefore); }); }); }); From f698808b43359fffd54855aa79e5ff4f40cdb40f Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 6 Aug 2020 14:00:48 -0700 Subject: [PATCH 18/22] . --- e2e/integration_tests/backends_test.ts | 26 +++++----- tfjs-backend-cpu/src/kernels/PadV2.ts | 4 +- tfjs-backend-cpu/src/kernels/Reshape.ts | 7 +-- .../src/kernels/SpaceToBatchND.ts | 35 +++++--------- .../src/kernels/SpaceToBatchND_test.ts | 48 ------------------- tfjs-backend-cpu/src/kernels/Transpose.ts | 45 +++++++++-------- 6 files changed, 57 insertions(+), 108 deletions(-) delete mode 100644 tfjs-backend-cpu/src/kernels/SpaceToBatchND_test.ts diff --git a/e2e/integration_tests/backends_test.ts b/e2e/integration_tests/backends_test.ts index f1a4704b2c4..858be68ea44 100644 --- a/e2e/integration_tests/backends_test.ts +++ b/e2e/integration_tests/backends_test.ts @@ -37,11 +37,11 @@ describe(`${SMOKE} backends`, () => { const webglBefore = tfc.engine().backend.numDataIds(); - // input is stored in webgl backend. tensorInfo refCount = 1. const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32'); + // input is stored in webgl backend. tensorInfo refCount = 1. - // input tensorInfo refCount = 2; const inputReshaped = tfc.reshape(input, [2, 2]); + // input tensorInfo refCount = 2; const webglAfter = tfc.engine().backend.numDataIds(); @@ -51,14 +51,14 @@ describe(`${SMOKE} backends`, () => { const cpuBefore = tfc.engine().backend.numDataIds(); - // input moved to cpu. tensorData refCount = 1. const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]); + // input moved to cpu. tensorData refCount = 1. const cpuAfter = tfc.engine().backend.numDataIds(); + expect(cpuAfter).toEqual(cpuBefore + 1); // input tensorData refCount = 3, reshape increase by 1, then output // tensor increase by 1. - expect(cpuAfter).toEqual(cpuBefore + 1); await tfc.setBackend('webgl'); @@ -67,20 +67,20 @@ describe(`${SMOKE} backends`, () => { await tfc.setBackend('cpu'); - // After dispose, tensorData refCount = 2. input.dispose(); + // After dispose, tensorData refCount = 2. // Input is not deleted, because refCount is 2. expect(tfc.engine().backend.numDataIds()).toEqual(cpuAfter); - // After dispose, tensorData refCount = 1. inputReshaped.dispose(); + // After dispose, tensorData refCount = 1. // Input is not deleted, because refCount is 1. expect(tfc.engine().backend.numDataIds()).toEqual(cpuAfter); - // After dispose, tensorData refCount = 0, data deleted. inputReshaped2.dispose(); + // After dispose, tensorData refCount = 0, data deleted. const after = tfc.engine().backend.numDataIds(); @@ -92,13 +92,13 @@ describe(`${SMOKE} backends`, () => { const cpuBefore = tfc.engine().backend.numDataIds(); + const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32'); // input is stored in cpu backend. tensorData refCount = 1, tensorInfo // refCount = 1. - const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32'); + const inputReshaped = tfc.reshape(input, [2, 2]); // input tensorData refCount = 3, reshape increase by 1, then output // tensor increase by 1. tensorInfo refCount = 1. - const inputReshaped = tfc.reshape(input, [2, 2]); const cpuAfter = tfc.engine().backend.numDataIds(); @@ -108,8 +108,8 @@ describe(`${SMOKE} backends`, () => { const webglBefore = tfc.engine().backend.numDataIds(); - // input moved to webgl. tensorInfo refCount = 3. const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]); + // input moved to webgl. tensorInfo refCount = 3. const webglAfter = tfc.engine().backend.numDataIds(); @@ -122,20 +122,20 @@ describe(`${SMOKE} backends`, () => { await tfc.setBackend('webgl'); - // After dispose, tensorInfo = 2. input.dispose(); + // After dispose, tensorInfo = 2. // Data is not deleted, because tensorInfo is 2. expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter); - // After dipose, tensorInfo = 1. inputReshaped.dispose(); + // After dipose, tensorInfo = 1. // Data is not deleted, because tensorInfo is 1. expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter); - // After dipose, tensorInfo = 1, data deleted. inputReshaped2.dispose(); + // After dipose, tensorInfo = 1, data deleted. const after = tfc.engine().backend.numDataIds(); diff --git a/tfjs-backend-cpu/src/kernels/PadV2.ts b/tfjs-backend-cpu/src/kernels/PadV2.ts index 731d0b381bd..d82462a8c86 100644 --- a/tfjs-backend-cpu/src/kernels/PadV2.ts +++ b/tfjs-backend-cpu/src/kernels/PadV2.ts @@ -1,6 +1,6 @@ /** * @license - * Copyright 2019 Google LLC. All Rights Reserved. + * 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 @@ -20,7 +20,7 @@ import {KernelConfig, KernelFunc, NumericDataType, PadV2, PadV2Attrs, PadV2Input import {MathBackendCPU} from '../backend_cpu'; import {assertNotComplex} from '../cpu_util'; -function padV2( +export function padV2( args: {inputs: PadV2Inputs, backend: MathBackendCPU, attrs: PadV2Attrs}): TensorInfo { const {inputs, backend, attrs} = args; diff --git a/tfjs-backend-cpu/src/kernels/Reshape.ts b/tfjs-backend-cpu/src/kernels/Reshape.ts index bc24b3384dc..e7d5e359266 100644 --- a/tfjs-backend-cpu/src/kernels/Reshape.ts +++ b/tfjs-backend-cpu/src/kernels/Reshape.ts @@ -15,13 +15,14 @@ * ============================================================================= */ -import {KernelConfig, KernelFunc, Reshape, ReshapeAttrs, ReshapeInputs} from '@tensorflow/tfjs-core'; +import {KernelConfig, KernelFunc, Reshape, ReshapeAttrs, ReshapeInputs, TensorInfo} from '@tensorflow/tfjs-core'; import {MathBackendCPU} from '../backend_cpu'; -function reshape( +export function reshape( args: - {inputs: ReshapeInputs, backend: MathBackendCPU, attrs: ReshapeAttrs}) { + {inputs: ReshapeInputs, backend: MathBackendCPU, attrs: ReshapeAttrs}): + TensorInfo { const {inputs, backend, attrs} = args; const {x} = inputs; const {shape} = attrs; diff --git a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts index 3553d6d0327..e59e2e4a91d 100644 --- a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts +++ b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts @@ -1,6 +1,6 @@ /** * @license - * Copyright 2019 Google LLC. All Rights Reserved. + * 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 @@ -15,16 +15,16 @@ * ============================================================================= */ -import {backend_util, KernelConfig, KernelFunc, NamedAttrMap, ReshapeAttrs, ReshapeInputs, SpaceToBatchND, SpaceToBatchNDAttrs, SpaceToBatchNDInputs, TensorInfo, TransposeAttrs, TransposeInputs} from '@tensorflow/tfjs-core'; +import {backend_util, KernelConfig, KernelFunc, ReshapeAttrs, ReshapeInputs, SpaceToBatchND, SpaceToBatchNDAttrs, SpaceToBatchNDInputs, TensorInfo, TransposeAttrs, TransposeInputs, util} from '@tensorflow/tfjs-core'; import {MathBackendCPU} from '../backend_cpu'; import {assertNotComplex} from '../cpu_util'; import {padV2Config} from './PadV2'; -import {reshapeConfig} from './Reshape'; -import {transposeConfig} from './Transpose'; +import {reshape} from './Reshape'; +import {transpose} from './Transpose'; -function spaceToBatchND(args: { +export function spaceToBatchND(args: { inputs: SpaceToBatchNDInputs, backend: MathBackendCPU, attrs: SpaceToBatchNDAttrs @@ -35,7 +35,7 @@ function spaceToBatchND(args: { assertNotComplex([x], 'spaceToBatchND'); - const prod = blockShape.reduce((a, b) => a * b); + const prod = util.sizeFromShape(blockShape); const completePaddings: Array<[number, number]> = [[0, 0]]; completePaddings.push(...(paddings as Array<[number, number]>)); @@ -61,34 +61,25 @@ function spaceToBatchND(args: { const reshapeInputs: ReshapeInputs = {x: paddedX}; const reshapeAttrs: ReshapeAttrs = {shape: reshapedPaddedShape}; - const paddedXReshaped = reshapeConfig.kernelFunc({ - inputs: reshapeInputs, - backend, - attrs: reshapeAttrs as {} as NamedAttrMap - }) as TensorInfo; + const paddedXReshaped = + reshape({inputs: reshapeInputs, backend, attrs: reshapeAttrs}); const transposeInputs: TransposeInputs = {x: paddedXReshaped}; const transposeAttrs: TransposeAttrs = {perm: permutedReshapedPaddedPermutation}; - const paddedXT = transposeConfig.kernelFunc({ - inputs: transposeInputs, - backend, - attrs: transposeAttrs as {} as NamedAttrMap - }) as TensorInfo; + const paddedXT = + transpose({inputs: transposeInputs, backend, attrs: transposeAttrs}); const resultReshapeInputs: ReshapeInputs = {x: paddedXT}; const resultReshapeAttrs: ReshapeAttrs = {shape: flattenShape}; - const result = reshapeConfig.kernelFunc({ - inputs: resultReshapeInputs, - backend, - attrs: resultReshapeAttrs as {} as NamedAttrMap - }); + const result = reshape( + {inputs: resultReshapeInputs, backend, attrs: resultReshapeAttrs}); backend.disposeDataSoft(paddedX.dataId); backend.disposeDataSoft(paddedXReshaped.dataId); backend.disposeDataSoft(paddedXT.dataId); - return result as TensorInfo; + return result; } export const spaceToBatchNDConfig: KernelConfig = { diff --git a/tfjs-backend-cpu/src/kernels/SpaceToBatchND_test.ts b/tfjs-backend-cpu/src/kernels/SpaceToBatchND_test.ts deleted file mode 100644 index f7766895f0d..00000000000 --- a/tfjs-backend-cpu/src/kernels/SpaceToBatchND_test.ts +++ /dev/null @@ -1,48 +0,0 @@ -/** - * @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 * as tf from '@tensorflow/tfjs-core'; -import {test_util} from '@tensorflow/tfjs-core'; - -const {expectArraysClose} = test_util; -// tslint:disable-next-line: no-imports-from-dist -import {describeWithFlags, ALL_ENVS} from '@tensorflow/tfjs-core/dist/jasmine_util'; - -describeWithFlags('SpaceToBatchND.', ALL_ENVS, () => { - it('has no memory leak.', async () => { - const initialDataIds = tf.engine().backend.numDataIds(); - - const t = tf.tensor4d([[[[1], [2]], [[3], [4]]]], [1, 2, 2, 1]); - const blockShape = [2, 2]; - const paddings = [[0, 0], [0, 0]]; - - const res = tf.spaceToBatchND(t, blockShape, paddings); - expect(res.shape).toEqual([4, 1, 1, 1]); - expectArraysClose(await res.data(), [1, 2, 3, 4]); - - const afterResDataIds = tf.engine().backend.numDataIds(); - // 1 input tensor and 1 result tensor - expect(afterResDataIds).toEqual(initialDataIds + 2); - - t.dispose(); - res.dispose(); - - const afterDisposeDataIds = tf.engine().backend.numDataIds(); - - expect(afterDisposeDataIds).toEqual(initialDataIds); - }); -}); diff --git a/tfjs-backend-cpu/src/kernels/Transpose.ts b/tfjs-backend-cpu/src/kernels/Transpose.ts index 05ae13f03cb..6935c118550 100644 --- a/tfjs-backend-cpu/src/kernels/Transpose.ts +++ b/tfjs-backend-cpu/src/kernels/Transpose.ts @@ -15,35 +15,40 @@ * ============================================================================= */ -import {KernelConfig, TypedArray} from '@tensorflow/tfjs-core'; +import {KernelConfig, KernelFunc, TensorInfo, Transpose, TransposeAttrs, TransposeInputs, TypedArray} from '@tensorflow/tfjs-core'; -import {Transpose, TransposeAttrs, TransposeInputs} from '@tensorflow/tfjs-core'; import {MathBackendCPU} from '../backend_cpu'; import {assertNotComplex} from '../cpu_util'; import {transposeImpl} from './Transpose_impl'; -export const transposeConfig: KernelConfig = { - kernelName: Transpose, - backendName: 'cpu', - kernelFunc: ({inputs, attrs, backend}) => { - const {x} = inputs as TransposeInputs; - const {perm} = attrs as {} as TransposeAttrs; - const cpuBackend = backend as MathBackendCPU; +export function transpose(args: { + inputs: TransposeInputs, + attrs: TransposeAttrs, + backend: MathBackendCPU +}): TensorInfo { + const {inputs, attrs, backend} = args; + const {x} = inputs; + const {perm} = attrs; - assertNotComplex(x, 'transpose'); + assertNotComplex(x, 'transpose'); - const xRank = x.shape.length; + const xRank = x.shape.length; - const newShape: number[] = new Array(xRank); - for (let i = 0; i < newShape.length; i++) { - newShape[i] = x.shape[perm[i]]; - } + const newShape: number[] = new Array(xRank); + for (let i = 0; i < newShape.length; i++) { + newShape[i] = x.shape[perm[i]]; + } - const values = cpuBackend.data.get(x.dataId).values as TypedArray; - const result = transposeImpl(values, x.shape, x.dtype, perm, newShape); + const values = backend.data.get(x.dataId).values as TypedArray; + const result = transposeImpl(values, x.shape, x.dtype, perm, newShape); - const dataId = cpuBackend.write(result, newShape, x.dtype); - return {dataId, shape: newShape, dtype: x.dtype}; - } + const dataId = backend.write(result, newShape, x.dtype); + return {dataId, shape: newShape, dtype: x.dtype}; +} + +export const transposeConfig: KernelConfig = { + kernelName: Transpose, + backendName: 'cpu', + kernelFunc: transpose as {} as KernelFunc }; From e8de89770a3aaaeddaf0f926c8cbb0130e56f79a Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 6 Aug 2020 14:36:48 -0700 Subject: [PATCH 19/22] . --- e2e/integration_tests/custom_op_test.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/e2e/integration_tests/custom_op_test.ts b/e2e/integration_tests/custom_op_test.ts index aa25905b625..b50b2d0675b 100644 --- a/e2e/integration_tests/custom_op_test.ts +++ b/e2e/integration_tests/custom_op_test.ts @@ -52,7 +52,6 @@ describe(`${SMOKE} custom model`, () => { ]); // A custom op that calls unmodularized kernels and modularized kernels. - // softmax ker tfconverter.registerOp('CustomOp', (nodeValue) => { const x = nodeValue.inputs[0]; const softMax = tfc.softmax(x); From 737ac90c7d26feedf6899e7d1b932d3ea3833b0d Mon Sep 17 00:00:00 2001 From: Na Li Date: Fri, 7 Aug 2020 18:19:42 -0700 Subject: [PATCH 20/22] Clean up. --- e2e/integration_tests/backends_test.ts | 46 ++------ e2e/integration_tests/custom_op_test.ts | 81 -------------- e2e/integration_tests/memory_leak_test.ts | 101 ++++++++++++++++++ tfjs-backend-cpu/src/backend_cpu.ts | 9 +- .../src/kernels/SpaceToBatchND.ts | 6 +- 5 files changed, 118 insertions(+), 125 deletions(-) delete mode 100644 e2e/integration_tests/custom_op_test.ts create mode 100644 e2e/integration_tests/memory_leak_test.ts diff --git a/e2e/integration_tests/backends_test.ts b/e2e/integration_tests/backends_test.ts index 858be68ea44..139ede3ea84 100644 --- a/e2e/integration_tests/backends_test.ts +++ b/e2e/integration_tests/backends_test.ts @@ -32,16 +32,14 @@ describe(`${SMOKE} backends`, () => { }); it(`from webgl to cpu.`, async () => { - // A backend with no refCounter. await tfc.setBackend('webgl'); const webglBefore = tfc.engine().backend.numDataIds(); const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32'); - // input is stored in webgl backend. tensorInfo refCount = 1. + // input is stored in webgl backend. const inputReshaped = tfc.reshape(input, [2, 2]); - // input tensorInfo refCount = 2; const webglAfter = tfc.engine().backend.numDataIds(); @@ -52,35 +50,23 @@ describe(`${SMOKE} backends`, () => { const cpuBefore = tfc.engine().backend.numDataIds(); const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]); - // input moved to cpu. tensorData refCount = 1. + // input moved to cpu. + + // Because input is moved to cpu, data should be deleted from webgl. + expect(tfc.findBackend('webgl').numDataIds()).toEqual(webglAfter - 1); const cpuAfter = tfc.engine().backend.numDataIds(); expect(cpuAfter).toEqual(cpuBefore + 1); - // input tensorData refCount = 3, reshape increase by 1, then output - // tensor increase by 1. - - await tfc.setBackend('webgl'); - - // Because input is moved to cpu, data should be deleted from webgl. - expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter - 1); - - await tfc.setBackend('cpu'); input.dispose(); - // After dispose, tensorData refCount = 2. - - // Input is not deleted, because refCount is 2. expect(tfc.engine().backend.numDataIds()).toEqual(cpuAfter); inputReshaped.dispose(); - // After dispose, tensorData refCount = 1. - // Input is not deleted, because refCount is 1. expect(tfc.engine().backend.numDataIds()).toEqual(cpuAfter); inputReshaped2.dispose(); - // After dispose, tensorData refCount = 0, data deleted. const after = tfc.engine().backend.numDataIds(); @@ -93,12 +79,9 @@ describe(`${SMOKE} backends`, () => { const cpuBefore = tfc.engine().backend.numDataIds(); const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32'); - // input is stored in cpu backend. tensorData refCount = 1, tensorInfo - // refCount = 1. + // input is stored in cpu backend. const inputReshaped = tfc.reshape(input, [2, 2]); - // input tensorData refCount = 3, reshape increase by 1, then output - // tensor increase by 1. tensorInfo refCount = 1. const cpuAfter = tfc.engine().backend.numDataIds(); @@ -109,33 +92,24 @@ describe(`${SMOKE} backends`, () => { const webglBefore = tfc.engine().backend.numDataIds(); const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]); - // input moved to webgl. tensorInfo refCount = 3. + // input moved to webgl. + + // Because input is moved to webgl, data should be deleted from cpu. + expect(tfc.findBackend('cpu').numDataIds()).toEqual(cpuAfter - 1); const webglAfter = tfc.engine().backend.numDataIds(); expect(webglAfter).toEqual(webglBefore + 1); - await tfc.setBackend('cpu'); - - // Because input is moved to webgl, data should be deleted from cpu. - expect(tfc.engine().backend.numDataIds()).toEqual(cpuAfter - 1); - - await tfc.setBackend('webgl'); - input.dispose(); - // After dispose, tensorInfo = 2. - // Data is not deleted, because tensorInfo is 2. expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter); inputReshaped.dispose(); - // After dipose, tensorInfo = 1. - // Data is not deleted, because tensorInfo is 1. expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter); inputReshaped2.dispose(); - // After dipose, tensorInfo = 1, data deleted. const after = tfc.engine().backend.numDataIds(); diff --git a/e2e/integration_tests/custom_op_test.ts b/e2e/integration_tests/custom_op_test.ts deleted file mode 100644 index b50b2d0675b..00000000000 --- a/e2e/integration_tests/custom_op_test.ts +++ /dev/null @@ -1,81 +0,0 @@ -import '@tensorflow/tfjs-backend-cpu'; - -import * as tfconverter from '@tensorflow/tfjs-converter'; -import * as tfc from '@tensorflow/tfjs-core'; - -import {SMOKE} from './constants'; - -const HOST = 'http://example.org'; -const MODEL_URL = `${HOST}/model.json`; - -const CUSTOM_OP_MODEL = { - node: [ - { - name: 'Input', - op: 'Placeholder', - attr: { - dtype: { - type: 1, // DT_FLOAT - }, - shape: {shape: {dim: [{size: 4}]}} - } - }, - {name: 'CustomOp', op: 'CustomOp', input: ['Input'], attr: {}} - ], - versions: {producer: 1.0, minConsumer: 3} -}; - -const weightsManifest: tfc.io.WeightsManifestEntry[] = - [{'name': 'Const', 'dtype': 'float32', 'shape': [1]}]; - -const bias = tfc.tensor1d([0], 'float32'); - -const CUSTOM_HTTP_MODEL_LOADER = { - load: async () => { - return { - modelTopology: CUSTOM_OP_MODEL, - weightSpecs: weightsManifest, - weightData: bias.dataSync(), - format: 'tfjs-graph-model', - generatedBy: '1.15', - convertedBy: '1.3.1' - }; - } -}; - -describe(`${SMOKE} custom model`, () => { - it('refCounter works correctly.', async () => { - const model = new tfconverter.GraphModel(MODEL_URL); - - spyOn(tfc.io, 'getLoadHandlers').and.returnValue([ - CUSTOM_HTTP_MODEL_LOADER - ]); - - // A custom op that calls unmodularized kernels and modularized kernels. - tfconverter.registerOp('CustomOp', (nodeValue) => { - const x = nodeValue.inputs[0]; - const softMax = tfc.softmax(x); - const clone = tfc.clone(softMax); - return [tfc.reshape(clone, [2, 2])]; - }); - - await model.load(); - - const before = tfc.memory().numTensors; - - const input = tfc.tensor1d([1, 2, 3, 4]); - const output = model.predict(input) as tfc.Tensor; - - tfc.test_util.expectArraysClose(await output.data(), [ - 0.032058604061603546, 0.08714432269334793, 0.23688283562660217, - 0.6439142823219299 - ]); - - input.dispose(); - output.dispose(); - - const after = tfc.memory().numTensors; - - expect(after).toEqual(before); - }); -}); diff --git a/e2e/integration_tests/memory_leak_test.ts b/e2e/integration_tests/memory_leak_test.ts new file mode 100644 index 00000000000..9ce0a2c96b0 --- /dev/null +++ b/e2e/integration_tests/memory_leak_test.ts @@ -0,0 +1,101 @@ +/** + * @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 '@tensorflow/tfjs-backend-cpu'; + +import * as tfconverter from '@tensorflow/tfjs-converter'; +import * as tfc from '@tensorflow/tfjs-core'; + +import {SMOKE} from './constants'; + +const HOST = 'http://example.org'; +const MODEL_URL = `${HOST}/model.json`; + +const CUSTOM_OP_MODEL = { + node: [ + { + name: 'Input', + op: 'Placeholder', + attr: { + dtype: { + type: 1, // DT_FLOAT + }, + shape: {shape: {dim: [{size: 4}]}} + } + }, + {name: 'CustomOp', op: 'CustomOp', input: ['Input'], attr: {}} + ], + versions: {producer: 1.0, minConsumer: 3} +}; + +const weightsManifest: tfc.io.WeightsManifestEntry[] = + [{'name': 'Const', 'dtype': 'float32', 'shape': [1]}]; + +const bias = tfc.tensor1d([0], 'float32'); + +const CUSTOM_HTTP_MODEL_LOADER = { + load: async () => { + return { + modelTopology: CUSTOM_OP_MODEL, + weightSpecs: weightsManifest, + weightData: bias.dataSync(), + format: 'tfjs-graph-model', + generatedBy: '1.15', + convertedBy: '1.3.1' + }; + } +}; + +describe( + `${SMOKE} A custom op that calls unmodularized kernels and modularized ` + + `kernels`, + () => { + it('should have no memory leak in a model run.', async () => { + const model = new tfconverter.GraphModel(MODEL_URL); + + spyOn(tfc.io, 'getLoadHandlers').and.returnValue([ + CUSTOM_HTTP_MODEL_LOADER + ]); + + // A custom op that calls unmodularized kernels and modularized kernels. + tfconverter.registerOp('CustomOp', (nodeValue) => { + const x = nodeValue.inputs[0]; + const softMax = tfc.softmax(x); + const clone = tfc.clone(softMax); + return [tfc.reshape(clone, [2, 2])]; + }); + + await model.load(); + + const before = tfc.memory().numTensors; + + const input = tfc.tensor1d([1, 2, 3, 4]); + const output = model.predict(input) as tfc.Tensor; + + tfc.test_util.expectArraysClose(await output.data(), [ + 0.032058604061603546, 0.08714432269334793, 0.23688283562660217, + 0.6439142823219299 + ]); + + input.dispose(); + output.dispose(); + + const after = tfc.memory().numTensors; + + expect(after).toEqual(before); + }); + }); diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index 0b15b45b037..a4c26676b2d 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -16,10 +16,7 @@ */ import * as tf from '@tensorflow/tfjs-core'; -import {engine, env} from '@tensorflow/tfjs-core'; -import {backend_util, buffer, slice_util, util} from '@tensorflow/tfjs-core'; -import {BackendTimingInfo, DataStorage, DataType, DataValues, KernelBackend, max, NumericDataType, Rank, Scalar, ShapeMap, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer, TypedArray, upcastType} from '@tensorflow/tfjs-core'; -import {kernel_impls} from '@tensorflow/tfjs-core'; +import {backend_util, BackendTimingInfo, buffer, DataStorage, DataType, DataValues, engine, env, kernel_impls, KernelBackend, max, NumericDataType, Rank, Scalar, ShapeMap, slice_util, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer, TensorInfo, TypedArray, upcastType, util} from '@tensorflow/tfjs-core'; const nonMaxSuppressionV3Impl = kernel_impls.nonMaxSuppressionV3Impl; const split = kernel_impls.split; @@ -174,7 +171,9 @@ export class MathBackendCPU extends KernelBackend { } } - disposeDataSoft(dataId: DataId): void { + disposeIntermediateTensorInfo(tensorInfo: TensorInfo): void { + const dataId = tensorInfo.dataId; + if (this.data.has(dataId)) { const tensorData = this.data.get(dataId); diff --git a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts index e59e2e4a91d..1e2796d4e44 100644 --- a/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts +++ b/tfjs-backend-cpu/src/kernels/SpaceToBatchND.ts @@ -75,9 +75,9 @@ export function spaceToBatchND(args: { const result = reshape( {inputs: resultReshapeInputs, backend, attrs: resultReshapeAttrs}); - backend.disposeDataSoft(paddedX.dataId); - backend.disposeDataSoft(paddedXReshaped.dataId); - backend.disposeDataSoft(paddedXT.dataId); + backend.disposeIntermediateTensorInfo(paddedX); + backend.disposeIntermediateTensorInfo(paddedXReshaped); + backend.disposeIntermediateTensorInfo(paddedXT); return result; } From 9d1cec2f974e046be5d2f8429e9c536b2f7a36bb Mon Sep 17 00:00:00 2001 From: Na Li Date: Fri, 7 Aug 2020 22:52:31 -0700 Subject: [PATCH 21/22] Remove decRef from engine. --- tfjs-core/src/engine.ts | 9 --------- 1 file changed, 9 deletions(-) diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index c0ea6bfe6c5..3412fa01b78 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -834,15 +834,6 @@ export class Engine implements TensorTracker, DataMover { this.state.tensorInfo.delete(a.dataId); } else { this.state.tensorInfo.get(a.dataId).refCount--; - - // Todo(linazhao): Remove decRef once reshape, clone and cast kernels - // are modularized. - // We added this to keep backend refCount to be in sync with engine - // refCount. - if (this.backendName === 'cpu') { - // tslint:disable-next-line: no-any - (info.backend as any).decRef(a.dataId); - } } // TODO(nsthorat): Construct an error and save the stack trace for // debugging when in debug mode. Creating a stack trace is too expensive From a77f655d893c65438fd58ae47fcf626e2106231a Mon Sep 17 00:00:00 2001 From: Na Li Date: Fri, 7 Aug 2020 23:00:17 -0700 Subject: [PATCH 22/22] . --- tfjs-backend-cpu/src/backend_cpu.ts | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index a4c26676b2d..62341737027 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -158,15 +158,11 @@ export class MathBackendCPU extends KernelBackend { disposeData(dataId: DataId): void { if (this.data.has(dataId)) { - const tensorData = this.data.get(dataId); - - if (tensorData.complexTensors != null) { - // Todo(linazhao): Change to disposeData once complex, real, and imag - // kernels are modularized and real and imag becomes `TensorInfo`. - tensorData.complexTensors.real.dispose(); - tensorData.complexTensors.imag.dispose(); + const {complexTensors} = this.data.get(dataId); + if (complexTensors != null) { + complexTensors.real.dispose(); + complexTensors.imag.dispose(); } - this.data.delete(dataId); } }