feat(webgl): update pool2d

Author: BenAnn

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

@@ -2,17 +2,24 @@
  * @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 }
 ) {
     const [stride_v = 1, stride_h = 1] = strides;
     const [padTop = 0, padLeft = 0] = paddings;
     const [ksize_x, ksize_y] = ksize;
+    const originShape = recoverShape(origin);
+
     return `
     // start函数
     void main(void) {
         float res = -1. / 0.;
+        int index = 0;
         // 获取output的坐标
         ivec4 out_pos = getOutputTensorPos();
         int b = out_pos[0];
@@ -40,10 +47,20 @@ function mainFunc(
                 }
                 // origin数据
                 float curr = getValueFromTensorPos_origin(out_pos[0], out_pos[1], oy, ox);
+                if (layer_run_time == 0 && ${global_pooling === true}) {
+                    if (curr > res) {
+                        index = ${originShape[2] * originShape[3]} * out_pos[1] + ${originShape[3]} * oy + ox;
+                    }
+                }
                 res = max(res, curr);
             }
         }
-        setOutput(res);
+        if (layer_run_time == 0 && ${global_pooling === true}) {
+            setOutput(float(index));
+        }
+        else {
+            setOutput(float(res));
+        }
     }
     `;
 }
@@ -52,14 +69,16 @@ export default {
     params: [
         'strides',
         'paddings',
-        'ksize'
+        'ksize',
+        'global_pooling'
     ],
     textureFuncConf: {
         origin: ['getValueFromTensorPos']
     },
     behaviors: [
         'isMax',
         'setPacked',
+        'setAdaptive',
         'isGlobalPooling'
     ]
 };

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}`;