set count_include_pad for avg_pool2d in TensorRT wrapper

Author: markrogersjr

src/contrib/subgraph/tensorrt_executor.cc

@@ -587,6 +587,17 @@ void AddPooling(
   } else {
     network->setPoolingOutputDimensionsFormula(nullptr);
   }
+  if (!is_global_pool) {
+    if (nodes[nid].attrs.count("count_include_pad")) {
+      if (nodes[nid].attrs.at("count_include_pad") == "True") {
+        pool_layer->setAverageCountExcludesPadding(false);
+      } else {
+       pool_layer->setAverageCountExcludesPadding(true);
+      }
+    } else {
+      pool_layer->setAverageCountExcludesPadding(true);
+    }
+  }
   nid2layer->emplace(nid, pool_layer);
 }
 

tests/python/tensorrt/test_avg_pool2d.py

@@ -0,0 +1,46 @@
+import numpy as np
+import mxnet as mx
+from mxnet import gluon
+import nnvm
+import tvm
+from tvm.contrib import graph_runtime
+
+
+def test_avg_pool2d():
+
+    # Generate the data
+    np.random.seed(0)
+    input_shape = [1, 1, 28, 28]
+    output_shape = [1, 10]
+    data = np.random.random(input_shape).astype('float32')
+    
+    # Baseline model in MXNet
+    net = gluon.nn.HybridSequential()
+    with net.name_scope():
+        net.add(gluon.nn.AvgPool2D(pool_size=3, strides=1, padding=1))
+        net.add(gluon.nn.Dense(10))
+    net.collect_params().initialize(mx.init.Xavier(), ctx=mx.cpu())
+    net.hybridize()
+    baseline_input = mx.nd.array(data, ctx=mx.cpu())
+    baseline_output = net(baseline_input).asnumpy()
+    
+    # Compiled model
+    sym, params = nnvm.frontend.from_mxnet(net)
+    target = tvm.target.cuda()
+    with nnvm.compiler.build_config(opt_level=3, ext_accel='tensorrt'):
+        graph, lib, params = nnvm.compiler.build(sym, target,
+                                                 shape={'data': input_shape},
+                                                 params=params)
+    compiled_model = graph_runtime.create(graph, lib, tvm.gpu())
+    compiled_input = tvm.nd.array(data, ctx=tvm.gpu())
+    compiled_model.set_input('data', compiled_input)
+    compiled_model.set_input(**params)
+    compiled_model.run()
+    compiled_output = compiled_model.get_output(0, tvm.nd.empty(output_shape)).asnumpy()
+    
+    # Compare outputs
+    np.testing.assert_almost_equal(baseline_output, compiled_output, decimal=3)
+
+
+if __name__ == '__main__':
+    test_avg_pool2d()