You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by pa...@apache.org on 2019/07/28 04:34:23 UTC
[incubator-mxnet] branch master updated: Fix subgraph with
custom_op (#15671)
This is an automated email from the ASF dual-hosted git repository.
patriczhao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new d464a47 Fix subgraph with custom_op (#15671)
d464a47 is described below
commit d464a47486cd4ec5470afbb291fbbe67b85044a0
Author: Zhennan Qin <zh...@intel.com>
AuthorDate: Sun Jul 28 12:34:00 2019 +0800
Fix subgraph with custom_op (#15671)
---
src/c_api/c_api_symbolic.cc | 4 ++-
src/c_api/c_api_test.cc | 4 ++-
src/executor/graph_executor.cc | 4 ++-
src/operator/subgraph/subgraph_property.h | 19 +++++++++++
tests/python/unittest/test_subgraph.py | 53 +++++++++++++++++++++++++++++++
5 files changed, 81 insertions(+), 3 deletions(-)
diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc
index df839b4..627acf4 100644
--- a/src/c_api/c_api_symbolic.cc
+++ b/src/c_api/c_api_symbolic.cc
@@ -1046,8 +1046,10 @@ int MXGenBackendSubgraph(SymbolHandle sym_handle, const char *backend_name,
for (auto property : subgraph_prop_list) {
nnvm::Graph g = Symbol2Graph(*s);
property->SetAttr("graph", g);
- g.attrs["subgraph_property"] = std::make_shared<nnvm::any>(std::move(property));
+ g.attrs["subgraph_property"] = std::make_shared<nnvm::any>(property);
g = ApplyPass(std::move(g), "BuildSubgraph");
+ property->RemoveAttr("graph");
+ g.attrs.erase("subgraph_property");
s->outputs = g.outputs;
}
*ret_sym_handle = s;
diff --git a/src/c_api/c_api_test.cc b/src/c_api/c_api_test.cc
index fff9269..a1f12fa 100644
--- a/src/c_api/c_api_test.cc
+++ b/src/c_api/c_api_test.cc
@@ -50,8 +50,10 @@ int MXBuildSubgraphByOpNames(SymbolHandle sym_handle,
g.outputs = s->outputs;
property->SetAttr("graph", g);
property->SetAttr("op_names", op_name_set);
- g.attrs["subgraph_property"] = std::make_shared<nnvm::any>(std::move(property));
+ g.attrs["subgraph_property"] = std::make_shared<nnvm::any>(property);
g = nnvm::ApplyPass(std::move(g), "BuildSubgraph");
+ property->RemoveAttr("graph");
+ g.attrs.erase("subgraph_property");
s->outputs = g.outputs;
}
}
diff --git a/src/executor/graph_executor.cc b/src/executor/graph_executor.cc
index 0ba85a9..c5da0a0 100644
--- a/src/executor/graph_executor.cc
+++ b/src/executor/graph_executor.cc
@@ -1688,8 +1688,10 @@ static nnvm::Symbol BuildSubgraph(const nnvm::Symbol& src, op::SubgraphPropertyP
g = InferForwardAttrs(g, arg_shapes, arg_dtypes, arg_stypes, default_ctx, ctx_map, in_arg_ctxes,
aux_state_ctxes, true);
subgraph_prop->SetAttr("graph", g);
- g.attrs["subgraph_property"] = std::make_shared<nnvm::any>(std::move(subgraph_prop));
+ g.attrs["subgraph_property"] = std::make_shared<nnvm::any>(subgraph_prop);
g = ApplyPass(std::move(g), "BuildSubgraph");
+ subgraph_prop->RemoveAttr("graph");
+ g.attrs.erase("subgraph_property");
ret.outputs = g.outputs;
return ret;
}
diff --git a/src/operator/subgraph/subgraph_property.h b/src/operator/subgraph/subgraph_property.h
index 0683e5e..e2c861e 100644
--- a/src/operator/subgraph/subgraph_property.h
+++ b/src/operator/subgraph/subgraph_property.h
@@ -331,6 +331,15 @@ class SubgraphProperty {
return it != attrs_.end();
}
/*!
+ * \brief Remove attr if the attr exists.
+ */
+ void RemoveAttr(const std::string& name) {
+ auto it = attrs_.find(name);
+ if (it != attrs_.end()) {
+ attrs_.erase(it);
+ }
+ }
+ /*!
* \brief Get the property type.
*/
SgPropertyType GetPropertyType() const { return type_; }
@@ -384,6 +393,16 @@ class SubgraphBackend {
return it != attrs_.end();
}
+ /*!
+ * \brief Remove attr if the attr exists.
+ */
+ void RemoveAttr(const std::string& name) {
+ auto it = attrs_.find(name);
+ if (it != attrs_.end()) {
+ attrs_.erase(it);
+ }
+ }
+
SubgraphPropertyPtr& RegisterSubgraphProperty(const SubgraphPropertyPtr prop) {
prop_ptr_.push_back(prop);
return prop_ptr_.back();
diff --git a/tests/python/unittest/test_subgraph.py b/tests/python/unittest/test_subgraph.py
index b5577d4..4c13f9c 100644
--- a/tests/python/unittest/test_subgraph.py
+++ b/tests/python/unittest/test_subgraph.py
@@ -144,6 +144,59 @@ def test_make_subgraph():
rtol=0.001, atol=0.0001)
+def test_subgraph_with_customOp():
+ class MyAdd(mx.operator.CustomOp):
+ def forward(self, is_train, req, in_data, out_data, aux):
+ self.assign(out_data[0], req[0], in_data[0] + 1)
+
+ def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
+ self.assign(in_grad[0], req[0], out_grad[0])
+
+ @mx.operator.register('MyAdd1')
+ class MyAdd1Prop(mx.operator.CustomOpProp):
+ def __init__(self):
+ super(MyAdd1Prop, self).__init__(need_top_grad=True)
+
+ def list_arguments(self):
+ return ['data']
+
+ def list_outputs(self):
+ return ['output']
+
+ def infer_shape(self, in_shape):
+ # inputs, outputs, aux
+ return [in_shape[0]], [in_shape[0]], []
+
+ def create_operator(self, ctx, shapes, dtypes):
+ return MyAdd()
+
+ @mx.operator.register('MyAdd2')
+ class MyAdd2Prop(mx.operator.CustomOpProp):
+ def __init__(self):
+ super(MyAdd2Prop, self).__init__(need_top_grad=True)
+
+ def list_arguments(self):
+ return ['data']
+
+ def list_outputs(self):
+ return ['output']
+
+ def infer_shape(self, in_shape):
+ # inputs, outputs, aux
+ return [in_shape[0]], [in_shape[0]], []
+
+ def create_operator(self, ctx, shapes, dtypes):
+ return MyAdd()
+
+ inp = mx.nd.zeros(shape=(100, 100))
+ a = mx.symbol.Variable('a')
+ b = a + 1
+ b = mx.symbol.Custom(data=a, op_type='MyAdd1')
+ c = mx.symbol.Custom(data=a, op_type='MyAdd2')
+ b.bind(mx.cpu(), {'a': inp}).forward()
+ c.bind(mx.cpu(), {'a': inp}).forward()
+ mx.nd.waitall()
+
if __name__ == '__main__':
import nose
nose.runmodule()