提交 971acff7 编写于 作者: T tensor-tang

add python interface for mkldnn_pool

上级 9d108a21
......@@ -2286,8 +2286,15 @@ class NormLayer(LayerBase):
@config_layer('pool')
class PoolLayer(LayerBase):
layer_type = 'pool'
def __init__(self, name, inputs, ceil_mode=True, **xargs):
super(PoolLayer, self).__init__(name, 'pool', 0, inputs=inputs, **xargs)
use_mkldnn = int(g_command_config_args.get("use_mkldnn", 0))
if self.layer_type == "mkldnn_pool":
config_assert(use_mkldnn, "mkldnn_pool only support MKLDNN")
self.layer_type = 'mkldnn_pool' if use_mkldnn else 'pool'
super(PoolLayer, self).__init__(
name, self.layer_type, 0, inputs=inputs, **xargs)
for input_index in xrange(len(self.inputs)):
input_layer = self.get_input_layer(input_index)
pool_conf = self.config.inputs[input_index].pool_conf
......@@ -2297,6 +2304,11 @@ class PoolLayer(LayerBase):
pool_conf.channels)
@config_layer('mkldnn_pool')
class MKLDNNPoolLayer(PoolLayer):
layer_type = 'mkldnn_pool'
@config_layer('pool3d')
class Pool3DLayer(LayerBase):
def __init__(self, name, inputs, ceil_mode=True, **xargs):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册