/**
|
* @license
|
* Copyright 2017 Google Inc. All Rights Reserved.
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
* you may not use this file except in compliance with the License.
|
* You may obtain a copy of the License at
|
*
|
* http://www.apache.org/licenses/LICENSE-2.0
|
*
|
* Unless required by applicable law or agreed to in writing, software
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
* See the License for the specific language governing permissions and
|
* limitations under the License.
|
* =============================================================================
|
*/
|
|
import {Conv2DInfo, Conv3DInfo} from '../../ops/conv_util';
|
|
import {GPGPUProgram} from './gpgpu_math';
|
|
export class AvgPool2DBackpropProgram implements GPGPUProgram {
|
variableNames = ['dy'];
|
outputShape: number[];
|
userCode: string;
|
|
constructor(convInfo: Conv2DInfo) {
|
this.outputShape = convInfo.inShape;
|
const filterHeight = convInfo.filterHeight;
|
const filterWidth = convInfo.filterWidth;
|
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 padTop = effectiveFilterHeight - 1 - convInfo.padInfo.top;
|
const padLeft = effectiveFilterWidth - 1 - convInfo.padInfo.left;
|
|
const avgMultiplier = 1 / (filterHeight * filterWidth);
|
|
this.userCode = `
|
const ivec2 pads = ivec2(${padTop}, ${padLeft});
|
const float avgMultiplier = float(${avgMultiplier});
|
|
void main() {
|
ivec4 coords = getOutputCoords();
|
int b = coords[0];
|
int d = coords[3];
|
|
ivec2 dyRCCorner = coords.yz - pads;
|
int dyRCorner = dyRCCorner.x;
|
int dyCCorner = dyRCCorner.y;
|
|
// Convolve dy(?, ?, d) with pos mask(:, :, d) to get dx(xR, xC, d).
|
// ? = to be determined. : = across all values in that axis.
|
float dotProd = 0.0;
|
for (int wR = 0; wR < ${effectiveFilterHeight};
|
wR += ${dilationHeight}) {
|
float dyR = float(dyRCorner + wR) / ${strideHeight}.0;
|
|
if (dyR < 0.0 || dyR >= ${convInfo.outHeight}.0 || fract(dyR) > 0.0) {
|
continue;
|
}
|
int idyR = int(dyR);
|
|
for (int wC = 0; wC < ${effectiveFilterWidth};
|
wC+= ${dilationWidth}) {
|
float dyC = float(dyCCorner + wC) / ${strideWidth}.0;
|
|
if (dyC < 0.0 || dyC >= ${convInfo.outWidth}.0 ||
|
fract(dyC) > 0.0) {
|
continue;
|
}
|
int idyC = int(dyC);
|
|
float dyValue = getDy(b, idyR, idyC, d);
|
|
dotProd += dyValue * avgMultiplier;
|
}
|
}
|
setOutput(dotProd);
|
}
|
`;
|
}
|
}
|
|
export class AvgPool3DBackpropProgram implements GPGPUProgram {
|
variableNames = ['dy'];
|
outputShape: number[];
|
userCode: string;
|
|
constructor(convInfo: Conv3DInfo) {
|
this.outputShape = convInfo.inShape;
|
const filterDepth = convInfo.filterDepth;
|
const filterHeight = convInfo.filterHeight;
|
const filterWidth = convInfo.filterWidth;
|
const strideDepth = convInfo.strideDepth;
|
const strideHeight = convInfo.strideHeight;
|
const strideWidth = convInfo.strideWidth;
|
const dilationDepth = convInfo.dilationDepth;
|
const dilationHeight = convInfo.dilationHeight;
|
const dilationWidth = convInfo.dilationWidth;
|
const effectiveFilterDepth = convInfo.effectiveFilterDepth;
|
const effectiveFilterHeight = convInfo.effectiveFilterHeight;
|
const effectiveFilterWidth = convInfo.effectiveFilterWidth;
|
|
const padFront = effectiveFilterDepth - 1 - convInfo.padInfo.front;
|
const padTop = effectiveFilterHeight - 1 - convInfo.padInfo.top;
|
const padLeft = effectiveFilterWidth - 1 - convInfo.padInfo.left;
|
|
const avgMultiplier = 1 / (filterDepth * filterHeight * filterWidth);
|
|
this.userCode = `
|
const ivec3 pads = ivec3(${padFront}, ${padTop}, ${padLeft});
|
const float avgMultiplier = float(${avgMultiplier});
|
|
void main() {
|
ivec5 coords = getOutputCoords();
|
int batch = coords.x;
|
int ch = coords.u;
|
|
ivec3 dyCorner = ivec3(coords.y, coords.z, coords.w) - pads;
|
int dyDCorner = dyCorner.x;
|
int dyRCorner = dyCorner.y;
|
int dyCCorner = dyCorner.z;
|
|
// Convolve dy(?, ?, ?, d) with pos mask(:, :, :, ch) to get
|
// dx(xD, xR, xC, ch).
|
// ? = to be determined. : = across all values in that axis.
|
float dotProd = 0.0;
|
|
for (int wD = 0; wD < ${effectiveFilterDepth};
|
wD += ${dilationDepth}) {
|
float dyD = float(dyDCorner + wD) / ${strideDepth}.0;
|
|
if (dyD < 0.0 || dyD >= ${convInfo.outDepth}.0 || fract(dyD) > 0.0) {
|
continue;
|
}
|
int idyD = int(dyD);
|
|
for (int wR = 0; wR < ${effectiveFilterHeight};
|
wR += ${dilationHeight}) {
|
float dyR = float(dyRCorner + wR) / ${strideHeight}.0;
|
|
if (dyR < 0.0 || dyR >= ${convInfo.outHeight}.0 ||
|
fract(dyR) > 0.0) {
|
continue;
|
}
|
int idyR = int(dyR);
|
|
for (int wC = 0; wC < ${effectiveFilterWidth};
|
wC += ${dilationWidth}) {
|
float dyC = float(dyCCorner + wC) / ${strideWidth}.0;
|
|
if (dyC < 0.0 || dyC >= ${convInfo.outWidth}.0 ||
|
fract(dyC) > 0.0) {
|
continue;
|
}
|
int idyC = int(dyC);
|
|
float dyValue = getDy(batch, idyD, idyR, idyC, ch);
|
|
dotProd += dyValue * avgMultiplier;
|
}
|
}
|
}
|
setOutput(dotProd);
|
}
|
`;
|
}
|
}
|