forked from tensorflow/tfjs
-
Notifications
You must be signed in to change notification settings - Fork 0
/
DepthwiseConv2dNativeBackpropInput.ts
110 lines (96 loc) · 4.1 KB
/
DepthwiseConv2dNativeBackpropInput.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
/**
* @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, DepthwiseConv2dNativeBackpropInput, DepthwiseConv2dNativeBackpropInputAttrs, DepthwiseConv2dNativeBackpropInputInputs, KernelConfig, KernelFunc, TensorBuffer, TensorInfo, TypedArray, util} from '@tensorflow/tfjs-core';
import {MathBackendCPU} from '../backend_cpu';
import {assertNotComplex} from '../cpu_util';
export function depthwiseConv2dNativeBackpropInput(args: {
inputs: DepthwiseConv2dNativeBackpropInputInputs,
backend: MathBackendCPU,
attrs: DepthwiseConv2dNativeBackpropInputAttrs
}): TensorInfo {
const {inputs, backend, attrs} = args;
const {dy, filter} = inputs;
const {strides, dilations, pad, dimRoundingMode, inputShape} = attrs;
assertNotComplex([dy, filter], 'depthwiseConv2DNativeBackpropInput');
const dyStrides = util.computeStrides(dy.shape);
const filterStrides = util.computeStrides(filter.shape);
const convInfo = backend_util.computeConv2DInfo(
inputShape, filter.shape as [number, number, number, number], strides,
dilations, pad, dimRoundingMode, true /* depthwise */);
const dx = new TensorBuffer(convInfo.inShape, 'float32');
const dxValues = dx.values;
const [dxS0, dxS1, dxS2] = dx.strides;
const dyValues = backend.data.get(dy.dataId).values as TypedArray;
const [dyS0, dyS1, dyS2] = dyStrides;
const fltValues = backend.data.get(filter.dataId).values as TypedArray;
const [fltS0, fltS1, fltS2] = filterStrides;
const {
batchSize,
filterHeight,
filterWidth,
inChannels,
inHeight,
inWidth,
outChannels,
outHeight,
outWidth,
strideHeight,
strideWidth
} = convInfo;
const topPad = filterHeight - 1 - convInfo.padInfo.top;
const leftPad = filterWidth - 1 - convInfo.padInfo.left;
const chMul = outChannels / inChannels;
for (let b = 0; b < batchSize; ++b) {
for (let d1 = 0; d1 < inChannels; ++d1) {
for (let xR = 0; xR < inHeight; ++xR) {
const xRCorner = xR - topPad;
const xRMin = Math.max(0, Math.ceil(xRCorner / strideHeight));
const yRMax =
Math.min(outHeight, (filterHeight + xRCorner) / strideHeight);
for (let xC = 0; xC < inWidth; ++xC) {
const xCCorner = xC - leftPad;
const xCMin = Math.max(0, Math.ceil(xCCorner / strideWidth));
const yCMax =
Math.min(outWidth, (filterWidth + xCCorner) / strideWidth);
let dotProd = 0;
for (let yR = xRMin; yR < yRMax; ++yR) {
const wR = yR * strideHeight - xRCorner;
for (let yC = xCMin; yC < yCMax; ++yC) {
const wC = yC * strideWidth - xCCorner;
const dyOffset = dyS0 * b + dyS1 * yR + dyS2 * yC;
const fltOffset = fltS0 * (filterHeight - 1 - wR) +
fltS1 * (filterWidth - 1 - wC) + fltS2 * d1;
for (let dm = 0; dm < chMul; ++dm) {
const d2 = d1 * chMul + dm;
const pixel = dyValues[dyOffset + d2];
const weight = fltValues[fltOffset + dm];
dotProd += pixel * weight;
}
}
}
dxValues[dxS0 * b + dxS1 * xR + dxS2 * xC + d1] = dotProd;
}
}
}
}
return backend.makeTensorInfo(dx.shape, dx.dtype, dx.values);
}
export const depthwiseConv2dNativeBackpropInputConfig: KernelConfig = {
kernelName: DepthwiseConv2dNativeBackpropInput,
backendName: 'cpu',
kernelFunc: depthwiseConv2dNativeBackpropInput as unknown as KernelFunc
};