diff --git a/tfjs-backend-cpu/src/backend_cpu.ts b/tfjs-backend-cpu/src/backend_cpu.ts index d9ecd0bb6be..eced73a2599 100644 --- a/tfjs-backend-cpu/src/backend_cpu.ts +++ b/tfjs-backend-cpu/src/backend_cpu.ts @@ -16,7 +16,7 @@ */ import * as tf 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'; +import {backend_util, BackendTimingInfo, 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; @@ -25,7 +25,6 @@ const topkImpl = kernel_impls.topkImpl; const whereImpl = kernel_impls.whereImpl; import * as seedrandom from 'seedrandom'; import {assertNotComplex} from './cpu_util'; -import {maxPoolPositions, pool} from './utils/pool_utils'; interface DataId {} @@ -1529,128 +1528,6 @@ export class MathBackendCPU extends KernelBackend { .slice(sliceBeginCoords, sliceSize) as T; } - maxPool(x: Tensor4D, convInfo: backend_util.Conv2DInfo): Tensor4D { - assertNotComplex(x, 'maxPool'); - const xValues = this.readSync(x.dataId) as TypedArray; - return pool(xValues, x.shape, x.dtype, x.strides, convInfo, 'max') - .toTensor() as Tensor4D; - } - - maxPoolBackprop( - dy: Tensor4D, x: Tensor4D, y: Tensor4D, - convInfo: backend_util.Conv2DInfo): Tensor4D { - assertNotComplex([x, y], 'maxPoolBackprop'); - - const xValues = this.readSync(x.dataId) as TypedArray; - const maxPosBuf = buffer( - convInfo.outShape, x.dtype, - maxPoolPositions(xValues, x.shape, x.dtype, convInfo).values); - const strideHeight = convInfo.strideHeight; - const strideWidth = convInfo.strideWidth; - const dilationHeight = convInfo.dilationHeight; - const dilationWidth = convInfo.dilationWidth; - const effectiveFilterHeight = convInfo.effectiveFilterHeight; - const effectiveFilterWidth = convInfo.effectiveFilterWidth; - const padLeft = effectiveFilterWidth - 1 - convInfo.padInfo.left; - const padTop = effectiveFilterHeight - 1 - convInfo.padInfo.top; - const dx = tf.buffer(x.shape, 'float32'); - - const dyBuf = this.bufferSync(dy); - - for (let b = 0; b < convInfo.batchSize; ++b) { - for (let d = 0; d < convInfo.inChannels; ++d) { - for (let dxR = 0; dxR < convInfo.inHeight; ++dxR) { - for (let dxC = 0; dxC < convInfo.inWidth; ++dxC) { - // Shader code begins. - const dyRCorner = dxR - padTop; - const dyCCorner = dxC - padLeft; - let dotProd = 0; - for (let wR = 0; wR < effectiveFilterHeight; wR += dilationHeight) { - const dyR = (dyRCorner + wR) / strideHeight; - if (dyR < 0 || dyR >= convInfo.outHeight || - Math.floor(dyR) !== dyR) { - continue; - } - for (let wC = 0; wC < effectiveFilterWidth; wC += dilationWidth) { - const dyC = (dyCCorner + wC) / strideWidth; - if (dyC < 0 || dyC >= convInfo.outWidth || - Math.floor(dyC) !== dyC) { - continue; - } - const maxPos = effectiveFilterHeight * effectiveFilterWidth - - 1 - (maxPosBuf.get(b, dyR, dyC, d) as number); - const curPos = wR * effectiveFilterWidth + wC; - - const mask = maxPos === curPos ? 1 : 0; - if (mask === 0) { - continue; - } - - const pixel = dyBuf.get(b, dyR, dyC, d); - dotProd += pixel * mask; - } - } - dx.set(dotProd, b, dxR, dxC, d); - } - } - } - } - return dx.toTensor(); - } - - avgPoolBackprop(dy: Tensor4D, x: Tensor4D, convInfo: backend_util.Conv2DInfo): - Tensor4D { - assertNotComplex([dy, x], 'avgPoolBackprop'); - - const strideHeight = convInfo.strideHeight; - const strideWidth = convInfo.strideWidth; - const filterHeight = convInfo.filterHeight; - const filterWidth = convInfo.filterWidth; - const dilationHeight = convInfo.dilationHeight; - const dilationWidth = convInfo.dilationWidth; - const effectiveFilterHeight = convInfo.effectiveFilterHeight; - const effectiveFilterWidth = convInfo.effectiveFilterWidth; - const padLeft = effectiveFilterWidth - 1 - convInfo.padInfo.left; - const padTop = effectiveFilterHeight - 1 - convInfo.padInfo.top; - const dx = tf.buffer(x.shape, 'float32'); - - const avgMultiplier = 1 / (filterHeight * filterWidth); - - const dyBuf = this.bufferSync(dy); - - for (let b = 0; b < convInfo.batchSize; ++b) { - for (let d = 0; d < convInfo.inChannels; ++d) { - for (let dxR = 0; dxR < convInfo.inHeight; ++dxR) { - for (let dxC = 0; dxC < convInfo.inWidth; ++dxC) { - // Shader code begins. - const dyRCorner = dxR - padTop; - const dyCCorner = dxC - padLeft; - let dotProd = 0; - for (let wR = 0; wR < effectiveFilterHeight; wR += dilationHeight) { - const dyR = (dyRCorner + wR) / strideHeight; - if (dyR < 0 || dyR >= convInfo.outHeight || - Math.floor(dyR) !== dyR) { - continue; - } - for (let wC = 0; wC < effectiveFilterWidth; wC += dilationWidth) { - const dyC = (dyCCorner + wC) / strideWidth; - if (dyC < 0 || dyC >= convInfo.outWidth || - Math.floor(dyC) !== dyC) { - continue; - } - - const pixel = dyBuf.get(b, dyR, dyC, d); - dotProd += pixel; - } - } - dx.set(dotProd * avgMultiplier, b, dxR, dxC, d); - } - } - } - } - return dx.toTensor(); - } - private pool3d( x: Tensor5D, convInfo: backend_util.Conv3DInfo, poolType: 'max'|'avg'): Tensor5D { @@ -2005,15 +1882,6 @@ export class MathBackendCPU extends KernelBackend { return dx.toTensor(); } - avgPool(x: Tensor4D, convInfo: backend_util.Conv2DInfo): Tensor4D { - assertNotComplex(x, 'avgPool'); - assertNotComplex(x, 'maxPool'); - const xValues = this.readSync(x.dataId) as TypedArray; - return pool(xValues, x.shape, x.dtype, x.strides, convInfo, 'avg') - .toTensor() - .toFloat() as Tensor4D; - } - resizeBilinear( x: Tensor4D, newHeight: number, newWidth: number, alignCorners: boolean): Tensor4D { diff --git a/tfjs-backend-cpu/src/kernels/AvgPool.ts b/tfjs-backend-cpu/src/kernels/AvgPool.ts new file mode 100644 index 00000000000..89b8523f5d1 --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/AvgPool.ts @@ -0,0 +1,61 @@ +/** + * @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 {AvgPool, AvgPoolAttrs, AvgPoolInputs, backend_util, KernelConfig, KernelFunc, TensorInfo, TypedArray, util} from '@tensorflow/tfjs-core'; + +import {MathBackendCPU} from '../backend_cpu'; +import {assertNotComplex} from '../cpu_util'; +import {pool} from '../utils/pool_utils'; +import {identity} from './Identity'; + +export function avgPool( + args: + {inputs: AvgPoolInputs, backend: MathBackendCPU, attrs: AvgPoolAttrs}): + TensorInfo { + const {inputs, backend, attrs} = args; + const {x} = inputs; + assertNotComplex(x, 'avgPool'); + const {filterSize, strides, pad, dimRoundingMode} = attrs; + const dilations = 1; + + util.assert( + backend_util.eitherStridesOrDilationsAreOne(strides, dilations), + () => 'Error in avgPool: Either strides or dilations must be 1. ' + + `Got strides ${strides} and dilations '${dilations}'`); + + const convInfo = backend_util.computePool2DInfo( + x.shape as [number, number, number, number], filterSize, strides, + dilations, pad, dimRoundingMode); + let res: TensorInfo; + + if (convInfo.filterWidth === 1 && convInfo.filterHeight === 1 && + util.arraysEqual(convInfo.inShape, convInfo.outShape)) { + res = identity({inputs: {x}, backend}); + } else { + const xValues = backend.data.get(x.dataId).values as TypedArray; + const strides = util.computeStrides(x.shape); + const buffer = pool(xValues, x.shape, x.dtype, strides, convInfo, 'avg'); + res = backend.makeTensorInfo( + convInfo.outShape, x.dtype, buffer.values as TypedArray); + } + return res; +} + +export const avgPoolConfig: KernelConfig = { + kernelName: AvgPool, + backendName: 'cpu', + kernelFunc: avgPool as {} as KernelFunc +}; diff --git a/tfjs-backend-cpu/src/kernels/AvgPoolBackprop.ts b/tfjs-backend-cpu/src/kernels/AvgPoolBackprop.ts new file mode 100644 index 00000000000..08e743f1061 --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/AvgPoolBackprop.ts @@ -0,0 +1,92 @@ +/** + * @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 {AvgPoolBackprop, AvgPoolBackpropAttrs, AvgPoolBackpropInputs, backend_util, buffer, KernelConfig, KernelFunc, Rank, TensorInfo} from '@tensorflow/tfjs-core'; + +import {MathBackendCPU} from '../backend_cpu'; +import {assertNotComplex} from '../cpu_util'; + +export function avgPoolBackprop(args: { + inputs: AvgPoolBackpropInputs, + backend: MathBackendCPU, + attrs: AvgPoolBackpropAttrs +}): TensorInfo { + const {inputs, backend, attrs} = args; + const {dy, input} = inputs; + const x = input; + assertNotComplex([dy, input], 'avgPoolBackprop'); + const {filterSize, strides, pad} = attrs; + + const convInfo = backend_util.computePool2DInfo( + x.shape as [number, number, number, number], filterSize, strides, + 1 /* dilations */, pad); + const strideHeight = convInfo.strideHeight; + const strideWidth = convInfo.strideWidth; + const filterHeight = convInfo.filterHeight; + const filterWidth = convInfo.filterWidth; + const dilationHeight = convInfo.dilationHeight; + const dilationWidth = convInfo.dilationWidth; + const effectiveFilterHeight = convInfo.effectiveFilterHeight; + const effectiveFilterWidth = convInfo.effectiveFilterWidth; + const padLeft = effectiveFilterWidth - 1 - convInfo.padInfo.left; + const padTop = effectiveFilterHeight - 1 - convInfo.padInfo.top; + const dx = + buffer(x.shape as [number, number, number, number], 'float32'); + + const avgMultiplier = 1 / (filterHeight * filterWidth); + + const dyData = backend.data.get(dy.dataId).values as Float32Array; + const dyBuf = buffer( + dy.shape as [number, number, number, number], 'float32', dyData); + + for (let b = 0; b < convInfo.batchSize; ++b) { + for (let d = 0; d < convInfo.inChannels; ++d) { + for (let dxR = 0; dxR < convInfo.inHeight; ++dxR) { + for (let dxC = 0; dxC < convInfo.inWidth; ++dxC) { + // Shader code begins. + const dyRCorner = dxR - padTop; + const dyCCorner = dxC - padLeft; + let dotProd = 0; + for (let wR = 0; wR < effectiveFilterHeight; wR += dilationHeight) { + const dyR = (dyRCorner + wR) / strideHeight; + if (dyR < 0 || dyR >= convInfo.outHeight || + Math.floor(dyR) !== dyR) { + continue; + } + for (let wC = 0; wC < effectiveFilterWidth; wC += dilationWidth) { + const dyC = (dyCCorner + wC) / strideWidth; + if (dyC < 0 || dyC >= convInfo.outWidth || + Math.floor(dyC) !== dyC) { + continue; + } + + const pixel = dyBuf.get(b, dyR, dyC, d); + dotProd += pixel; + } + } + dx.set(dotProd * avgMultiplier, b, dxR, dxC, d); + } + } + } + } + return backend.makeTensorInfo(dx.shape, dx.dtype, dx.values); +} + +export const avgPoolBackpropConfig: KernelConfig = { + kernelName: AvgPoolBackprop, + backendName: 'cpu', + kernelFunc: avgPoolBackprop as {} as KernelFunc +}; diff --git a/tfjs-backend-cpu/src/kernels/MaxPool.ts b/tfjs-backend-cpu/src/kernels/MaxPool.ts new file mode 100644 index 00000000000..7c7cc43d54e --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/MaxPool.ts @@ -0,0 +1,61 @@ +/** + * @license + * Copyright 2020 Google LLC. All Rights Reserved. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * ============================================================================= + */ +import {backend_util, KernelConfig, KernelFunc, MaxPool, MaxPoolAttrs, MaxPoolInputs, TensorInfo, TypedArray, util} from '@tensorflow/tfjs-core'; + +import {MathBackendCPU} from '../backend_cpu'; +import {assertNotComplex} from '../cpu_util'; +import {pool} from '../utils/pool_utils'; +import {identity} from './Identity'; + +export function maxPool( + args: + {inputs: MaxPoolInputs, backend: MathBackendCPU, attrs: MaxPoolAttrs}): + TensorInfo { + const {inputs, backend, attrs} = args; + const {x} = inputs; + assertNotComplex(x, 'maxPool'); + const {filterSize, strides, pad, dimRoundingMode} = attrs; + const dilations = 1; + + util.assert( + backend_util.eitherStridesOrDilationsAreOne(strides, dilations), + () => 'Error in maxPool: Either strides or dilations must be 1. ' + + `Got strides ${strides} and dilations '${dilations}'`); + + const convInfo = backend_util.computePool2DInfo( + x.shape as [number, number, number, number], filterSize, strides, + dilations, pad, dimRoundingMode); + let res: TensorInfo; + + if (convInfo.filterWidth === 1 && convInfo.filterHeight === 1 && + util.arraysEqual(convInfo.inShape, convInfo.outShape)) { + res = identity({inputs: {x}, backend}); + } else { + const xValues = backend.data.get(x.dataId).values as TypedArray; + const strides = util.computeStrides(x.shape); + const buffer = pool(xValues, x.shape, x.dtype, strides, convInfo, 'max'); + res = backend.makeTensorInfo( + convInfo.outShape, x.dtype, buffer.values as TypedArray); + } + return res; +} + +export const maxPoolConfig: KernelConfig = { + kernelName: MaxPool, + backendName: 'cpu', + kernelFunc: maxPool as {} as KernelFunc +}; diff --git a/tfjs-backend-cpu/src/kernels/MaxPoolBackprop.ts b/tfjs-backend-cpu/src/kernels/MaxPoolBackprop.ts new file mode 100644 index 00000000000..69dc207faa1 --- /dev/null +++ b/tfjs-backend-cpu/src/kernels/MaxPoolBackprop.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 {backend_util, buffer, KernelConfig, KernelFunc, MaxPoolBackprop, MaxPoolBackpropAttrs, MaxPoolBackpropInputs, Rank, TensorInfo, TypedArray} from '@tensorflow/tfjs-core'; + +import {MathBackendCPU} from '../backend_cpu'; +import {assertNotComplex} from '../cpu_util'; +import {maxPoolPositions} from '../utils/pool_utils'; + +export function maxPoolBackprop(args: { + inputs: MaxPoolBackpropInputs, + backend: MathBackendCPU, + attrs: MaxPoolBackpropAttrs +}): TensorInfo { + const {inputs, backend, attrs} = args; + const {dy, input, output} = inputs; + const x = input; + assertNotComplex([input, output], 'maxPoolBackprop'); + const {filterSize, strides, pad, dimRoundingMode} = attrs; + + const convInfo = backend_util.computePool2DInfo( + x.shape as [number, number, number, number], filterSize, strides, + 1 /* dilations */, pad, dimRoundingMode); + const xValues = backend.data.get(x.dataId).values as TypedArray; + const maxPosBuf = buffer( + convInfo.outShape, x.dtype, + maxPoolPositions(xValues, x.shape, x.dtype, convInfo).values); + const strideHeight = convInfo.strideHeight; + const strideWidth = convInfo.strideWidth; + const dilationHeight = convInfo.dilationHeight; + const dilationWidth = convInfo.dilationWidth; + const effectiveFilterHeight = convInfo.effectiveFilterHeight; + const effectiveFilterWidth = convInfo.effectiveFilterWidth; + const padLeft = effectiveFilterWidth - 1 - convInfo.padInfo.left; + const padTop = effectiveFilterHeight - 1 - convInfo.padInfo.top; + const dx = + buffer(x.shape as [number, number, number, number], 'float32'); + + const dyData = backend.data.get(dy.dataId).values as Float32Array; + const dyBuf = buffer( + dy.shape as [number, number, number, number], 'float32', dyData); + + for (let b = 0; b < convInfo.batchSize; ++b) { + for (let d = 0; d < convInfo.inChannels; ++d) { + for (let dxR = 0; dxR < convInfo.inHeight; ++dxR) { + for (let dxC = 0; dxC < convInfo.inWidth; ++dxC) { + // Shader code begins. + const dyRCorner = dxR - padTop; + const dyCCorner = dxC - padLeft; + let dotProd = 0; + for (let wR = 0; wR < effectiveFilterHeight; wR += dilationHeight) { + const dyR = (dyRCorner + wR) / strideHeight; + if (dyR < 0 || dyR >= convInfo.outHeight || + Math.floor(dyR) !== dyR) { + continue; + } + for (let wC = 0; wC < effectiveFilterWidth; wC += dilationWidth) { + const dyC = (dyCCorner + wC) / strideWidth; + if (dyC < 0 || dyC >= convInfo.outWidth || + Math.floor(dyC) !== dyC) { + continue; + } + const maxPos = effectiveFilterHeight * effectiveFilterWidth - 1 - + (maxPosBuf.get(b, dyR, dyC, d) as number); + const curPos = wR * effectiveFilterWidth + wC; + + const mask = maxPos === curPos ? 1 : 0; + if (mask === 0) { + continue; + } + + const pixel = dyBuf.get(b, dyR, dyC, d); + dotProd += pixel * mask; + } + } + dx.set(dotProd, b, dxR, dxC, d); + } + } + } + } + return backend.makeTensorInfo(dx.shape, dx.dtype, dx.values); +} + +export const maxPoolBackpropConfig: KernelConfig = { + kernelName: MaxPoolBackprop, + backendName: 'cpu', + kernelFunc: maxPoolBackprop as {} as KernelFunc +}; diff --git a/tfjs-backend-cpu/src/register_all_kernels.ts b/tfjs-backend-cpu/src/register_all_kernels.ts index 499a271523b..bec5c0cd270 100644 --- a/tfjs-backend-cpu/src/register_all_kernels.ts +++ b/tfjs-backend-cpu/src/register_all_kernels.ts @@ -27,6 +27,8 @@ import {asinConfig} from './kernels/Asin'; import {asinhConfig} from './kernels/Asinh'; import {atanConfig} from './kernels/Atan'; import {atanhConfig} from './kernels/Atanh'; +import {avgPoolConfig} from './kernels/AvgPool'; +import {avgPoolBackpropConfig} from './kernels/AvgPoolBackprop'; import {batchNormConfig} from './kernels/BatchNorm'; import {castConfig} from './kernels/Cast'; import {ceilConfig} from './kernels/Ceil'; @@ -56,6 +58,8 @@ import {logConfig} from './kernels/Log'; import {log1pConfig} from './kernels/Log1p'; import {logicalNotConfig} from './kernels/LogicalNot'; import {maxConfig} from './kernels/Max'; +import {maxPoolConfig} from './kernels/MaxPool'; +import {maxPoolBackpropConfig} from './kernels/MaxPoolBackprop'; import {maxPoolWithArgmaxConfig} from './kernels/MaxPoolWithArgmax'; import {multiplyConfig} from './kernels/Multiply'; import {nonMaxSuppressionV4Config} from './kernels/NonMaxSuppressionV4'; @@ -95,6 +99,8 @@ const kernelConfigs: KernelConfig[] = [ asinhConfig, atanConfig, atanhConfig, + avgPoolConfig, + avgPoolBackpropConfig, batchNormConfig, castConfig, ceilConfig, @@ -123,6 +129,8 @@ const kernelConfigs: KernelConfig[] = [ logConfig, log1pConfig, logicalNotConfig, + maxPoolConfig, + maxPoolBackpropConfig, maxPoolWithArgmaxConfig, maxConfig, multiplyConfig,