Skip to content

Commit

Permalink
[OP] Conv2d and Depthwise Conv2d for Raspberry Pi (apache#49)
Browse files Browse the repository at this point in the history
* [TUTORIAL] ImageNet Inference on Raspberry Pi

* Update tvm
  • Loading branch information
ZihengJiang authored and tqchen committed May 29, 2018
1 parent 0aafbff commit f8db4af
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 2 deletions.
3 changes: 2 additions & 1 deletion nnvm/python/nnvm/testing/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ def create_workload(net, batch_size, image_shape=(3, 224, 224),
params : dict of str to NDArray
The parameters.
"""
image_shape = (3, 224, 224)
if image_shape is None:
image_shape = (3, 224, 224)
data_shape = (batch_size,) + image_shape
params = {}
g = graph.create(net)
Expand Down
8 changes: 7 additions & 1 deletion nnvm/python/nnvm/top/nn.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ def compute_conv2d(attrs, inputs, _):
assert layout == "NCHW", "only support nchw for now"
assert dilation == (1, 1), "not support dilate now"
if groups == 1:
out = topi.nn.conv2d_nchw(inputs[0], inputs[1], strides, padding)
out = topi.nn.conv2d(inputs[0], inputs[1], strides, padding)
elif groups == get_const_int(inputs[0].shape[1]) and groups == channels:
out = topi.nn.depthwise_conv2d_nchw(inputs[0], inputs[1], strides, padding)
else:
Expand All @@ -128,6 +128,12 @@ def schedule_conv2d(attrs, outs, target):
return topi.cuda.schedule_conv2d_nchw(outs)
return topi.cuda.schedule_depthwise_conv2d_nchw(outs)
# naive schedule

if tvm.target.current_target() == tvm.target.rasp():
if groups == 1:
return topi.rasp.schedule_conv2d(outs)
return topi.rasp.schedule_depthwise_conv2d(outs)

return tvm.create_schedule([x.op for x in outs])

reg.register_pattern("conv2d", OpPattern.OUT_ELEMWISE_FUSABLE)
Expand Down

0 comments on commit f8db4af

Please sign in to comment.