Merge pull request #460 from JingyuanZhang/master

feat(webgl): add new op pool2d_avg_adaptive

Author: zhongkai

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

@@ -10,6 +10,7 @@ import conv2d_depthwise from './shader/conv2d_depthwise';
 import depthwise_conv2d from './shader/depthwise_conv2d';
 import conv2d_elementwise_add from './shader/conv2d_elementwise_add';
 import pool2d from './shader/pool2d';
+import pool2d_avg_adaptive from './shader/pool2d_avg_adaptive';
 import pool2d_max from './shader/pool2d_max';
 import pool2d_winograd from './shader/pool2d_winograd';
 import elementwise_add from './shader/elementwise_add';
@@ -117,6 +118,7 @@ const ops = {
     rnn_cell,
     rnn_origin,
     pool2d_avg,
+    pool2d_avg_adaptive,
     prelu: dynamic('prelu'),
     relu6: dynamic('relu6'),
     leakyRelu: dynamic('leakyRelu'),

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

@@ -64,6 +64,7 @@ export default {
         origin: ['getValueFromTensorPos']
     },
     behaviors: [
+        'isAdaptiveAvg',
         'isMax',
         'setPacked',
         'setAdaptive',

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

@@ -0,0 +1,45 @@
+/**
+ * @file pool2d_avg_adaptive
+ * @desc https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/nn/AdaptiveAvgPool2D_cn.html#adaptiveavgpool2d
+ */
+
+function mainFunc(
+    { origin },
+    { strides = [], paddings = [], ksize }
+) {
+    const [stride_v = 1, stride_h = 1] = strides;
+    const [padTop = 0, padLeft = 0] = paddings;
+    const [ksize_y, ksize_x] = ksize;
+    const H = origin.height_shape;
+    const W = origin.width_shape;
+    return `
+    // start函数
+    void main(void) {
+        float res = 0.0;
+        // 获取output的坐标
+        ivec4 out_pos = getOutputTensorPos();
+        int i = out_pos[2] * ${stride_v} - ${padTop};
+        int j = out_pos[3] * ${stride_h} - ${padLeft};
+        int hstart = int(floor(float(i) * float(${H}) / float(${ksize_y})));
+        int hend = int(ceil(float(i + 1) * float(${H}) / float(${ksize_y})));
+        int wstart = int(floor(float(j) * float(${W}) / float(${ksize_x})));
+        int wend = int(ceil(float(j + 1) * float(${W}) / float(${ksize_x})));
+        for (int fy = hstart; fy < hend; fy++) {
+            for (int fx = wstart; fx < wend; fx++) {
+                float curr = getValueFromTensorPos_origin(out_pos[0], out_pos[1], fy, fx);
+                res += curr;
+            }
+        }
+        int count_pool = (hend - hstart) * (wend - wstart);
+        res = res / float(count_pool);
+        setOutput(res);
+    }
+    `;
+}
+export default {
+    mainFunc,
+    textureFuncConf: {
+        origin: ['getValueFromTensorPos']
+    },
+    behaviors: []
+};

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

@@ -92,6 +92,17 @@ const behaviors : Behaviors = {
         }
     },
 
+    isAdaptiveAvg() {
+        const {
+            adaptive,
+            pooling_type
+        } = this.processedAttrs;
+
+        if (adaptive && pooling_type === 'avg') {
+            this.name += '_avg_adaptive';
+        }
+    },
+
     isMax() {
         const type = this.processedAttrs['pooling_type'] === 'max' ? 1 : 0;
         this.processedAttrs['pooling_type'] = type;