Merge pull request #141 from BenAnn/master

feat(webgl): update pool2d

Author: yueshuangyan

packages/paddlejs-backend-webgl/src/ops/shader/pool2d_max.ts

@@ -2,17 +2,34 @@
  * @file pool2d_max
  */
 
+function recoverShape({ total_shape, channel, height_shape, width_shape }) {
+    return [total_shape / channel / height_shape / width_shape, channel, height_shape, width_shape];
+}
+
 function mainFunc(
     { origin },
-    { strides = [], paddings = [], ksize }
+    { strides = [], paddings = [], ksize, global_pooling, runtime }
 ) {
     const [stride_v = 1, stride_h = 1] = strides;
     const [padTop = 0, padLeft = 0] = paddings;
     const [ksize_x, ksize_y] = ksize;
+    const originShape = recoverShape(origin);
+    let computedIndex = '';
+    let outputCode = 'setOutput(float(res));';
+    if (runtime === 0 && global_pooling === true) {
+        computedIndex = `
+            if (curr > res) {
+                index = ${originShape[2] * originShape[3]} * out_pos[1] + ${originShape[3]} * oy + ox;
+            }
+        `;
+        outputCode = 'setOutput(float(index));';
+    }
+
     return `
     // start函数
     void main(void) {
         float res = -1. / 0.;
+        int index = 0;
         // 获取output的坐标
         ivec4 out_pos = getOutputTensorPos();
         int b = out_pos[0];
@@ -40,10 +57,11 @@ function mainFunc(
                 }
                 // origin数据
                 float curr = getValueFromTensorPos_origin(out_pos[0], out_pos[1], oy, ox);
+                ${computedIndex}
                 res = max(res, curr);
             }
         }
-        setOutput(res);
+        ${outputCode}
     }
     `;
 }
@@ -52,14 +70,17 @@ export default {
     params: [
         'strides',
         'paddings',
-        'ksize'
+        'ksize',
+        'global_pooling',
+        'runtime'
     ],
     textureFuncConf: {
         origin: ['getValueFromTensorPos']
     },
     behaviors: [
         'isMax',
         'setPacked',
+        'setAdaptive',
         'isGlobalPooling'
     ]
 };

packages/paddlejs-backend-webgl/src/ops/utils.ts

@@ -21,7 +21,7 @@ const inputParams = [
     'limit',
     'channel',
     'total_shape',
-    'numbers_shape',
+    'numbers_shape'
 ];
 
 const outParams = [
@@ -44,7 +44,7 @@ const baseParams = {
     ]
 };
 
-function getTensorParams(inputTensors: Tensor[], ownParams: [], fShaderParams: object): opInfo {
+function getTensorParams(inputTensors: Tensor[], ownParams: [], fShaderParams: object, runtime: number): opInfo {
     const tensorsParams = {};
     const opParams = {};
     const tensorNames = [] as string[];
@@ -97,6 +97,9 @@ function getTensorParams(inputTensors: Tensor[], ownParams: [], fShaderParams: o
     if (fShaderParams['active_function']) {
         opParams['active_function'] = fShaderParams['active_function'];
     }
+
+    opParams['runtime'] = runtime;
+
     return { textureParams: tensorsParams, opParams, active_function: fShaderParams['active_function'] };
 }
 

packages/paddlejs-backend-webgl/src/webgl/buildShader.ts

@@ -20,7 +20,9 @@ export default function buildShader(textureConf, type, inputTensors, fShaderPara
         const { params = {}, mainFunc, textureFuncConf = {}, commonFuncConf } = ops[opName];
 
         // textureList: [filter, origin, bias]
-        const { textureParams, opParams, active_function } = getTensorParams(inputTensors, params, fShaderParams);
+        const { textureParams, opParams, active_function } = getTensorParams(
+            inputTensors, params, fShaderParams, runtime
+        );
 
         const prefixCode = genPrefixCode(textureConf);
 

packages/paddlejs-backend-webgl/test/op/data/max_pool2d_with_index.json

@@ -0,0 +1,106 @@
+{
+    "chunkNum": 0,
+    "ops": [
+        {
+            "attrs": {
+                "op_device": ""
+            },
+            "inputs": {
+                "X": [
+                    "feed"
+                ]
+            },
+            "outputs": {
+                "Out": [
+                    "tmp_43"
+                ]
+            },
+            "type": "feed"
+        }, 
+        {
+            "attrs": {
+                "adaptive": true,
+                "global_pooling": false,
+                "ksize": [
+                    1, 
+                    1
+                ],
+                "op_device": "",
+                "paddings": [
+                    0,
+                    0
+                ], 
+                "strides": [
+                    1,
+                    1
+                ]
+            },
+            "inputs": {
+                "X": [
+                    "tmp_43"
+                ]
+            },
+            "outputs": {
+                "Mask": [
+                    "max_pool2d_with_index_0.tmp_1"
+                ],
+                "Out": [
+                    "max_pool2d_with_index_0.tmp_0"
+                ]
+            },
+            "type": "max_pool2d_with_index"
+        },
+        {
+            "attrs": {
+                "op_device": ""
+            }, 
+            "inputs": {
+                "X": [
+                    "max_pool2d_with_index_0.tmp_0"
+                ]
+            },
+            "outputs": {
+                "Out": [
+                    "fetch"
+                ]
+            },
+            "type": "fetch"
+        }
+    ],
+    "vars": [
+        {
+            "name": "tmp_43", 
+            "shape": [
+                1,
+                3, 
+                3, 
+                3
+            ],
+            "data": [
+                1, 2, 3, 1, 1, 5, 1, 7, 1,
+                9, 10, 1, 1, 1, 8, 1, 1, 1,
+                1, 7, 1, 9, 14, 1, 1, 1, 8
+            ]
+        },
+        {
+            "name": "max_pool2d_with_index_0.tmp_0",
+            "shape": [
+                1, 
+                3, 
+                1, 
+                1
+            ],
+            "data": [7, 10, 14]
+        },
+        {
+            "name": "max_pool2d_with_index_0.tmp_1",
+            "shape": [
+                1, 
+                3, 
+                1, 
+                1
+            ],
+            "data": [7, 10, 22]
+        }
+    ]
+}

packages/paddlejs-core/src/opFactory/opBehaviors.ts

@@ -20,6 +20,18 @@ const behaviors : Behaviors = {
         }
     },
 
+    setAdaptive() {
+        if (
+            this.attrs.adaptive
+            && this.attrs.ksize.length === 2
+            && this.attrs.ksize[0] === 1
+            && this.attrs.ksize[1] === 1
+        ) {
+            this.attrs.adaptive = false;
+            this.attrs.global_pooling = true;
+        }
+    },
+
     isGlobalPooling() {
         const counter = this.input.X[0] || {};
         const length = (counter.shape && counter.shape.length) || 0;

packages/paddlejs-core/src/opFactory/opDataBuilder.ts

@@ -127,7 +127,8 @@ export default class OpData {
             scale: 'scale',
             bias: 'bias',
             mean: 'mean',
-            variance: 'variance'
+            variance: 'variance',
+            mask: 'out'
         };
 
 
@@ -170,6 +171,10 @@ export default class OpData {
                 this.name = 'reshape2';
             }
 
+            if (this.name.indexOf('max_pool2d_with_index') > -1) {
+                this.name = 'pool2d_max';
+            }
+
             const tensorData: ModelVar[] = this.tensorData;
             // unique behavior
             const opKey = `${GLOBALS.backend}_${this.name}`;