From 3bc227b36c6ca963e39f922188f39d799c0ad091 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 20 Mar 2018 01:48:03 +0000 Subject: [PATCH 01/71] Test input a graph. --- src/operator/nn/control_flow.cc | 62 +++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 src/operator/nn/control_flow.cc diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc new file mode 100644 index 000000000000..47c66c37b20b --- /dev/null +++ b/src/operator/nn/control_flow.cc @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include +#include +#include +#include +#include +#include +#include + +namespace mxnet { +namespace op { + +static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs) { + CHECK(attrs.g != nullptr); +} + +NNVM_REGISTER_OP(Foreach) +.set_num_inputs(3) +.set_num_outputs(1) +.set_attr("FListInputNames", + [](const NodeAttrs& attrs) { + return std::vector{"fn", "data1", "data2"}; +}) +.set_attr("FInputGraph", + [](const NodeAttrs& attrs) { + return 0; +}) +//.set_attr("FInferShape", ConvolutionShape) +//.set_attr("FInferType", ConvolutionType) +.describe(R"code(test)code" ADD_FILELINE) +//.set_attr_parser(ParamParser) +//.set_attr("FInferStorageType", ActivationStorageType) +.set_attr("FComputeEx", ForeachComputeExCPU) +.add_argument("fn", "Symbol", "Input graph.") +.add_argument("data1", "NDArray-or-Symbol", "Input1.") +.add_argument("data2", "NDArray-or-Symbol", "Input2."); +//.add_arguments(ActivationParam::__FIELDS__()); + +} // namespace op +} // namespace mxnet From dfdd28e0ba8dccb0d9015dbc53c9b7d56f92ec58 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 21 Mar 2018 12:05:25 -0700 Subject: [PATCH 02/71] Update foreach to execute the subgraph. --- include/mxnet/imperative.h | 22 +-- src/operator/nn/control_flow.cc | 237 +++++++++++++++++++++++++++++++- 2 files changed, 242 insertions(+), 17 deletions(-) diff --git a/include/mxnet/imperative.h b/include/mxnet/imperative.h index 758ce8513213..c82c7afeb49b 100644 --- a/include/mxnet/imperative.h +++ b/include/mxnet/imperative.h @@ -183,18 +183,18 @@ class Imperative { std::vector* p_save_inputs = nullptr, std::vector* p_save_outputs = nullptr); /*! \brief */ - OpStatePtr Invoke(const Context& default_ctx, - const nnvm::NodeAttrs& attrs, - const std::vector& inputs, - const std::vector& outputs); + static OpStatePtr Invoke(const Context& default_ctx, + const nnvm::NodeAttrs& attrs, + const std::vector& inputs, + const std::vector& outputs); /*! \brief */ - OpStatePtr InvokeOp(const Context& ctx, - const nnvm::NodeAttrs& attrs, - const std::vector& inputs, - const std::vector& outputs, - const std::vector& req, - const DispatchMode dispatch_mode, - OpStatePtr state = OpStatePtr()); + static OpStatePtr InvokeOp(const Context& ctx, + const nnvm::NodeAttrs& attrs, + const std::vector& inputs, + const std::vector& outputs, + const std::vector& req, + const DispatchMode dispatch_mode, + OpStatePtr state = OpStatePtr()); /*! \brief mark variables for computing gradients. */ void MarkVariables(const std::vector& variables, const std::vector& grad_reqs, diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 47c66c37b20b..15a375d7f9b7 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -24,19 +24,247 @@ #include #include #include +#include "../operator_common.h" +#include "../../imperative/imperative_utils.h" namespace mxnet { namespace op { +void RunGraph(const nnvm::IndexedGraph& idx, + const std::vector arrays, + size_t node_start, size_t node_end, + std::vector&& array_reqs, + std::vector&& ref_count, + std::vector *p_states, + const DispatchModeVector &dispatch_modes) { + using namespace nnvm; + using namespace imperative; + static auto& createop = nnvm::Op::GetAttr("FCreateOpState"); + static auto& is_layer_backward = Op::GetAttr("TIsLayerOpBackward"); + + std::vector& states = *p_states; + std::vector ndinputs, ndoutputs; + ShapeVector arg_shapes; + DTypeVector arg_dtypes; + std::vector req; + + for (size_t i = node_start; i < node_end; ++i) { + const nnvm::IndexedGraph::Node& node = idx[i]; + if (node.source->op() == nullptr) continue; + auto num_outputs = node.source->num_outputs(); + ndinputs.clear(); + ndinputs.reserve(node.inputs.size()); + for (const auto& j : node.inputs) { + ndinputs.emplace_back(arrays[idx.entry_id(j)]); + CHECK(!ndinputs.back()->is_none()) << idx[j.node_id].source->attrs.name << " " << j.index; + } + ndoutputs.clear(); + ndoutputs.reserve(num_outputs); + req.clear(); + req.reserve(num_outputs); + for (size_t j = 0; j < num_outputs; ++j) { + size_t eid = idx.entry_id(i, j); + ndoutputs.emplace_back(arrays[eid]); + req.push_back(array_reqs[eid]); + CHECK(!ndoutputs.back()->is_none()); + } + const Context& ctx = ndoutputs[0]->ctx(); + const DispatchMode dispatch_mode = dispatch_modes[i]; + if (createop.count(node.source->op())) { + arg_shapes.clear(); + arg_dtypes.clear(); + arg_shapes.reserve(ndinputs.size()); + arg_dtypes.reserve(ndinputs.size()); + for (size_t i = 0; i < ndinputs.size(); ++i) { + arg_shapes.emplace_back(ndinputs[i]->shape()); + arg_dtypes.emplace_back(ndinputs[i]->dtype()); + } + states[i] = createop[node.source->op()]( + node.source->attrs, ctx, arg_shapes, arg_dtypes); + Imperative::InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, req, + dispatch_mode, states[i]); + } else if (is_layer_backward.get(node.source->op(), false)) { + nnvm::Node* fwd_node = node.source->control_deps[0].get(); + auto fwd_node_id = idx.node_id(fwd_node); + Imperative::InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, + req, dispatch_mode, states[fwd_node_id]); + } else { + Imperative::InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, + req, dispatch_mode); + } + } +} + +static void ExecSubgraph(nnvm::Graph &g, const OpContext& ctx, + const std::vector& cinputs, + const std::vector& req, + const std::vector& coutputs) { + using namespace nnvm; + using namespace imperative; + const auto& idx = g.indexed_graph(); + size_t num_inputs = idx.input_nodes().size(); + + CHECK_EQ(num_inputs, cinputs.size()) + << "The subgraph requires " << num_inputs << " but got " << cinputs.size(); + + Context default_ctx = cinputs[0].ctx(); + for (size_t i = 0; i < cinputs.size(); ++i) { + CHECK_EQ(cinputs[i].ctx(), default_ctx) + << "The subgraph requires all inputs to live on the same context. But " + << idx[idx.input_nodes()[0]].source->attrs.name << " is on " << default_ctx + << " while " << idx[idx.input_nodes()[i]].source->attrs.name << " is on " + << cinputs[i].ctx(); + } + + // TODO(zhengda) we might want to buffer them. + std::vector buff; + std::vector states; + std::vector inputs = cinputs; + std::vector outputs = coutputs; + + // Allocate entries + states.resize(idx.num_nodes()); + buff.resize(idx.num_node_entries()); + states.reserve(idx.num_nodes()); + std::vector arrays; + arrays.reserve(buff.size()); + for (size_t i = 0; i < buff.size(); ++i) arrays.push_back(&buff[i]); + for (size_t i = 0; i < num_inputs; ++i) { + arrays[idx.entry_id(idx.input_nodes()[i], 0)] = &inputs[i]; + } + for (size_t i = 0; i < idx.outputs().size(); ++i) { + auto eid = idx.entry_id(idx.outputs()[i]); + if (!arrays[eid]->is_none()) outputs[i] = arrays[eid]->Detach(); + arrays[eid] = &outputs[i]; + } + + // Allocate memory for the NDArrays + std::vector ref_count = g.GetAttr >( + ctx.is_train ? "full_ref_count" : "forward_ref_count"); + + std::vector array_reqs(arrays.size(), kWriteTo); + for (size_t i = 0; i < idx.num_node_entries(); ++i) { + if (ref_count[i] == 0) array_reqs[i] = kNullOp; + } + + const auto& mem_plan = g.GetAttr( + ctx.is_train ? "full_mem_plan" : "forward_mem_plan"); + AllocateMemory(g, idx, default_ctx, 0, idx.num_node_entries(), + mem_plan, arrays, &array_reqs); + + const auto& dispatch_modes = g.GetAttr("dispatch_mode"); + + RunGraph(idx, arrays, 0, idx.num_nodes(), std::move(array_reqs), + std::move(ref_count), &states, dispatch_modes); +} + static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, const OpContext& ctx, const std::vector& inputs, const std::vector& req, const std::vector& outputs) { CHECK(attrs.g != nullptr); + nnvm::Graph &g = *attrs.g; + + printf("test1\n"); + // If this is inference, we only need the forward memory plan. + bool has_mem_plan = !ctx.is_train && g.attrs.count("forward_mem_plan"); + printf("test2\n"); + // If this is training, we need the full memory plan. + has_mem_plan = has_mem_plan || (ctx.is_train && g.attrs.count("full_mem_plan")); + printf("test3\n"); + // If we don't have a memory plan yet, we need to create a memory plan. + if (!has_mem_plan) { + const auto& idx = g.indexed_graph(); + nnvm::StorageVector storage(idx.num_node_entries(), exec::kBadStorageID); + for (const auto i : idx.input_nodes()) + storage[idx.entry_id(i, 0)] = exec::kExternalStorageID; + printf("test4\n"); + const auto& stypes = g.GetAttr("storage_type"); + printf("test5\n"); + CHECK_EQ(stypes.size(), storage.size()); + for (size_t i = 0; i < stypes.size(); i++) { + if (stypes[i] != kDefaultStorage) + storage[i] = exec::kDynamicStorageID; + } + + auto mem_plan = imperative::PlanMemory( + &g, std::move(storage), g.GetAttr >( + ctx.is_train ? "full_ref_count" : "forward_ref_count")); + printf("test6\n"); + // TODO(zhengda) we need to be careful of changing graph attributes. + // It's not thread-safe. + g.attrs[ctx.is_train ? "full_mem_plan" : "forward_mem_plan"] + = std::make_shared(std::move(mem_plan)); + printf("test7\n"); + } + printf("test8\n"); + ExecSubgraph(g, ctx, inputs, req, outputs); +} + +static bool ForeachShape(const nnvm::NodeAttrs& attrs, + std::vector *in_shape, + std::vector *out_shape) { + nnvm::ShapeVector shape_inputs = *in_shape; + auto g = attrs.g; + CHECK(g); + // TODO(zhengda) This can also be called in the execution engine. + // We need to make it thread-safe. + imperative::CheckAndInferShape(g.get(), std::move(shape_inputs), true); + const auto& shapes = g->GetAttr("shape"); + CHECK(g->outputs.size() == 1); + uint32_t eid = g->indexed_graph().entry_id(g->outputs[0]); + (*out_shape)[0] = shapes[eid]; + return true; +} + +static bool ForeachType(const nnvm::NodeAttrs& attrs, + std::vector *in_type, std::vector *out_type) { + nnvm::DTypeVector dtype_inputs = *in_type; + auto g = attrs.g; + CHECK(g); + // TODO(zhengda) This can also be called in the execution engine. + // We need to make it thread-safe. + imperative::CheckAndInferType(g.get(), std::move(dtype_inputs), true); + const auto &dtypes = g->GetAttr("dtype"); + CHECK(g->outputs.size() == 1); + uint32_t eid = g->indexed_graph().entry_id(g->outputs[0]); + (*out_type)[0] = dtypes[eid]; + return true; +} + +static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, + const int dev_mask, + DispatchMode* dispatch_mode, + std::vector *in_attrs, + std::vector *out_attrs) { + auto g = attrs.g; + CHECK(g); + printf("test1\n"); + const auto& idx = g->indexed_graph(); + CHECK(idx.input_nodes().size() == in_attrs->size()); + exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); + StorageTypeVector &storage_type_inputs = *in_attrs; + printf("test2\n"); + imperative::CheckAndInferStorageType(g.get(), std::move(dev_masks), + std::move(storage_type_inputs), true); + printf("test3\n"); + *dispatch_mode = DispatchMode::kFComputeEx; + const auto& stypes = g->GetAttr("storage_type"); + auto &outputs = idx.outputs(); + CHECK(outputs.size() == out_attrs->size()); + printf("test4\n"); + for (size_t i = 0; i < out_attrs->size(); i++) { + (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; + } + printf("test5\n"); + return true; } NNVM_REGISTER_OP(Foreach) +.describe(R"code(Foreach)code" ADD_FILELINE) +//.set_attr_parser(ParamParser) +.set_attr("FInferStorageType", ForeachStorageType) .set_num_inputs(3) .set_num_outputs(1) .set_attr("FListInputNames", @@ -47,16 +275,13 @@ NNVM_REGISTER_OP(Foreach) [](const NodeAttrs& attrs) { return 0; }) -//.set_attr("FInferShape", ConvolutionShape) -//.set_attr("FInferType", ConvolutionType) -.describe(R"code(test)code" ADD_FILELINE) -//.set_attr_parser(ParamParser) -//.set_attr("FInferStorageType", ActivationStorageType) +.set_attr("FInferShape", ForeachShape) +.set_attr("FInferType", ForeachType) .set_attr("FComputeEx", ForeachComputeExCPU) .add_argument("fn", "Symbol", "Input graph.") .add_argument("data1", "NDArray-or-Symbol", "Input1.") .add_argument("data2", "NDArray-or-Symbol", "Input2."); -//.add_arguments(ActivationParam::__FIELDS__()); +//.add_arguments(ForeachParam::__FIELDS__()); } // namespace op } // namespace mxnet From 42c735d96012f39c07f903310a30d949db7004dc Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 21 Mar 2018 21:55:27 +0000 Subject: [PATCH 03/71] print inputs/outputs in foreach. --- src/operator/nn/control_flow.cc | 38 +++++++++++++++++++++------------ 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 15a375d7f9b7..3abe6baa50ce 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -95,6 +95,21 @@ void RunGraph(const nnvm::IndexedGraph& idx, } } +static inline void print_dims(const mxnet::NDArray &arr, const std::string name = "") { + printf("%s: ", name.c_str()); + for (size_t i = 0; i < arr.shape().ndim(); i++) + printf("%ld, ", arr.shape()[i]); + printf("\n"); +} + +static inline void print(const mxnet::NDArray &arr, const std::string name = "") { + print_dims(arr, name); + float *data = (float *) arr.data().dptr_; + for (size_t i = 0; i < arr.shape().Size(); i++) + printf("%f, ", data[i]); + printf("\n"); +} + static void ExecSubgraph(nnvm::Graph &g, const OpContext& ctx, const std::vector& cinputs, const std::vector& req, @@ -151,11 +166,19 @@ static void ExecSubgraph(nnvm::Graph &g, const OpContext& ctx, ctx.is_train ? "full_mem_plan" : "forward_mem_plan"); AllocateMemory(g, idx, default_ctx, 0, idx.num_node_entries(), mem_plan, arrays, &array_reqs); + print(inputs[0], "data1"); + print(inputs[1], "data2"); + print(outputs[0], "output"); const auto& dispatch_modes = g.GetAttr("dispatch_mode"); RunGraph(idx, arrays, 0, idx.num_nodes(), std::move(array_reqs), std::move(ref_count), &states, dispatch_modes); + + printf("After running graph\n"); + print(inputs[0], "data1"); + print(inputs[1], "data2"); + print(outputs[0], "output"); } static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, @@ -166,22 +189,17 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, CHECK(attrs.g != nullptr); nnvm::Graph &g = *attrs.g; - printf("test1\n"); // If this is inference, we only need the forward memory plan. bool has_mem_plan = !ctx.is_train && g.attrs.count("forward_mem_plan"); - printf("test2\n"); // If this is training, we need the full memory plan. has_mem_plan = has_mem_plan || (ctx.is_train && g.attrs.count("full_mem_plan")); - printf("test3\n"); // If we don't have a memory plan yet, we need to create a memory plan. if (!has_mem_plan) { const auto& idx = g.indexed_graph(); nnvm::StorageVector storage(idx.num_node_entries(), exec::kBadStorageID); for (const auto i : idx.input_nodes()) storage[idx.entry_id(i, 0)] = exec::kExternalStorageID; - printf("test4\n"); const auto& stypes = g.GetAttr("storage_type"); - printf("test5\n"); CHECK_EQ(stypes.size(), storage.size()); for (size_t i = 0; i < stypes.size(); i++) { if (stypes[i] != kDefaultStorage) @@ -191,14 +209,11 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, auto mem_plan = imperative::PlanMemory( &g, std::move(storage), g.GetAttr >( ctx.is_train ? "full_ref_count" : "forward_ref_count")); - printf("test6\n"); // TODO(zhengda) we need to be careful of changing graph attributes. // It's not thread-safe. g.attrs[ctx.is_train ? "full_mem_plan" : "forward_mem_plan"] = std::make_shared(std::move(mem_plan)); - printf("test7\n"); } - printf("test8\n"); ExecSubgraph(g, ctx, inputs, req, outputs); } @@ -240,24 +255,19 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, std::vector *out_attrs) { auto g = attrs.g; CHECK(g); - printf("test1\n"); const auto& idx = g->indexed_graph(); CHECK(idx.input_nodes().size() == in_attrs->size()); exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); - StorageTypeVector &storage_type_inputs = *in_attrs; - printf("test2\n"); + StorageTypeVector storage_type_inputs = *in_attrs; imperative::CheckAndInferStorageType(g.get(), std::move(dev_masks), std::move(storage_type_inputs), true); - printf("test3\n"); *dispatch_mode = DispatchMode::kFComputeEx; const auto& stypes = g->GetAttr("storage_type"); auto &outputs = idx.outputs(); CHECK(outputs.size() == out_attrs->size()); - printf("test4\n"); for (size_t i = 0; i < out_attrs->size(); i++) { (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; } - printf("test5\n"); return true; } From a4d6a645707f1e32c71441bb3aa12fa69c10c343 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 23 Mar 2018 22:08:53 +0000 Subject: [PATCH 04/71] Remove print. --- src/operator/nn/control_flow.cc | 27 ++------------------------- 1 file changed, 2 insertions(+), 25 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 3abe6baa50ce..614de336e938 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -56,7 +56,8 @@ void RunGraph(const nnvm::IndexedGraph& idx, ndinputs.reserve(node.inputs.size()); for (const auto& j : node.inputs) { ndinputs.emplace_back(arrays[idx.entry_id(j)]); - CHECK(!ndinputs.back()->is_none()) << idx[j.node_id].source->attrs.name << " " << j.index; + CHECK(!ndinputs.back()->is_none()) << idx[j.node_id].source->attrs.name + << " " << j.index; } ndoutputs.clear(); ndoutputs.reserve(num_outputs); @@ -95,21 +96,6 @@ void RunGraph(const nnvm::IndexedGraph& idx, } } -static inline void print_dims(const mxnet::NDArray &arr, const std::string name = "") { - printf("%s: ", name.c_str()); - for (size_t i = 0; i < arr.shape().ndim(); i++) - printf("%ld, ", arr.shape()[i]); - printf("\n"); -} - -static inline void print(const mxnet::NDArray &arr, const std::string name = "") { - print_dims(arr, name); - float *data = (float *) arr.data().dptr_; - for (size_t i = 0; i < arr.shape().Size(); i++) - printf("%f, ", data[i]); - printf("\n"); -} - static void ExecSubgraph(nnvm::Graph &g, const OpContext& ctx, const std::vector& cinputs, const std::vector& req, @@ -166,19 +152,10 @@ static void ExecSubgraph(nnvm::Graph &g, const OpContext& ctx, ctx.is_train ? "full_mem_plan" : "forward_mem_plan"); AllocateMemory(g, idx, default_ctx, 0, idx.num_node_entries(), mem_plan, arrays, &array_reqs); - print(inputs[0], "data1"); - print(inputs[1], "data2"); - print(outputs[0], "output"); const auto& dispatch_modes = g.GetAttr("dispatch_mode"); - RunGraph(idx, arrays, 0, idx.num_nodes(), std::move(array_reqs), std::move(ref_count), &states, dispatch_modes); - - printf("After running graph\n"); - print(inputs[0], "data1"); - print(inputs[1], "data2"); - print(outputs[0], "output"); } static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, From 037caa0f5aff33f2d7420b03447da01d093f9c44 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 23 Mar 2018 22:09:23 +0000 Subject: [PATCH 05/71] add test code for foreach. --- tests/python/unittest/test_operator.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index e7976e01f9d8..70cb1ecddc2e 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5663,6 +5663,24 @@ def test_float16_min_max(): assert np.finfo('float16').max == mx.nd.max(a).asscalar() +@with_seed() +def test_foreach(): + v1 = mx.sym.var("v1") + v2 = mx.sym.var("v2") + v3 = mx.sym.var("v3") + v4 = mx.sym.var("v4") + g = v1 + v2 + op = mx.sym.Foreach(g, v3, v4) + arr1 = mx.nd.random.uniform(shape=(5, 2)) + arr2 = mx.nd.random.uniform(shape=(5, 2)) + e = op.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}) + e.forward() + for y in e.outputs: + y.wait_to_read() + print(y) + print(arr1 + arr2) + + @with_seed() def test_squeeze_op(): def check_squeeze_op(shape, axis=None): From 6e4b9fb52b7b4d29252bd7f8a1352b720bf5365c Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Thu, 5 Apr 2018 21:37:27 +0000 Subject: [PATCH 06/71] exec foreach outside the engine. --- src/executor/attach_op_execs_pass.cc | 8 ++++++++ src/executor/exec_pass.h | 3 +++ src/executor/graph_executor.cc | 24 +++++++++++++++++++++--- tests/python/unittest/test_operator.py | 5 +++-- 4 files changed, 35 insertions(+), 5 deletions(-) diff --git a/src/executor/attach_op_execs_pass.cc b/src/executor/attach_op_execs_pass.cc index 697e4869a049..b78149975113 100644 --- a/src/executor/attach_op_execs_pass.cc +++ b/src/executor/attach_op_execs_pass.cc @@ -201,6 +201,10 @@ class FComputeExecutor : public StorageFallbackOpExecutor { return exec_type_; } + bool HasSubgraph() const override { + return attrs_.g != nullptr; + } + explicit FComputeExecutor(const NodeAttrs& attrs, FCompute fcompute, ExecType exec_type, const std::vector &mutate_idx) : StorageFallbackOpExecutor(mutate_idx), @@ -226,6 +230,10 @@ class FComputeExExecutor : public OpExecutor { void Setup() override {} + bool HasSubgraph() const override { + return attrs_.g != nullptr; + } + ExecType exec_type() const override { return exec_type_; } diff --git a/src/executor/exec_pass.h b/src/executor/exec_pass.h index 99b1b162eaee..aa28e37500d5 100644 --- a/src/executor/exec_pass.h +++ b/src/executor/exec_pass.h @@ -64,6 +64,9 @@ class OpExecutor { OpContext op_ctx; /*! \brief virtual destructor */ virtual ~OpExecutor() {} + virtual bool HasSubgraph() const { + return false; + } /*! * \brief Setup the executor for given NDArray member * this can be called multiple times if NDArray changed during reshape. diff --git a/src/executor/graph_executor.cc b/src/executor/graph_executor.cc index 7a15f6c931c7..fc668b733639 100644 --- a/src/executor/graph_executor.cc +++ b/src/executor/graph_executor.cc @@ -1378,7 +1378,11 @@ void GraphExecutor::BulkTrainingOpSegs(size_t total_num_nodes) { // check if the segment relies on external input, or exceeds maxinum number of node, // or requires async ops if (node->is_variable() || nid - topo_start > num_nodes_threshold || - op_node.exec->exec_type() != ExecType::kSync) { + op_node.exec->exec_type() != ExecType::kSync || + // If the node has a subgraph, we shouldn't add it to the segment. + // We'll execute the node separately from other nodes. + // CreateCachedSegOpr creates a segment excluding nodes with subgraphs. + op_node.exec->HasSubgraph()) { // create a new segment for the previous nodes if the current one cannot be bulked cached_seg_opr_[topo_start] = this->CreateCachedSegOpr(topo_start, nid); topo_start = nid + 1; @@ -1403,7 +1407,11 @@ void GraphExecutor::BulkTrainingOpSegs(size_t total_num_nodes) { continue; } if (idx[nid].source->is_variable() || nid - topo_start > num_nodes_threshold || - op_node.exec->exec_type() != ExecType::kSync) { + op_node.exec->exec_type() != ExecType::kSync || + // If the node has a subgraph, we shouldn't add it to the segment. + // We'll execute the node separately from other nodes. + // CreateCachedSegOpr creates a segment excluding nodes with subgraphs. + op_node.exec->HasSubgraph()) { cached_seg_opr_[topo_start] = this->CreateCachedSegOpr(topo_start, nid); topo_start = nid + 1; } else { @@ -1437,7 +1445,11 @@ void GraphExecutor::BulkInferenceOpSegs() { // Variables do not need to be segmented at inference time. if (node->is_variable()) continue; - if (op_node.exec->exec_type() != ExecType::kSync) { + if (op_node.exec->exec_type() != ExecType::kSync || + // If the node has a subgraph, we shouldn't add it to the segment. + // We'll execute the node separately from other nodes. + // CreateCachedSegOpr creates a segment excluding nodes with subgraphs. + op_node.exec->HasSubgraph()) { cached_seg_opr_[topo_start] = this->CreateCachedSegOpr(topo_start, nid); topo_start = nid + 1; } @@ -1503,6 +1515,9 @@ void GraphExecutor::RunOps(bool is_train, size_t topo_start, size_t topo_end) { CHECK_EQ(opnode.exec->in_array.size(), 1U); CHECK_EQ(opnode.exec->out_array.size(), 1U); CopyFromTo(opnode.exec->in_array[0], &(opnode.exec->out_array[0])); + } else if (opnode.exec->HasSubgraph()) { + // If the node contains a subgraph, we can't execute it in the engine. + opnode.exec->Run(opnode.exec->op_ctx.run_ctx, false); } else if (opnode.cached_opr != nullptr) { bool profiling = profiler::Profiler::Get()->GetState() == profiler::Profiler::kRunning; Engine::Get()->Push(opnode.cached_opr, opnode.ctx, 0, profiling); @@ -1537,6 +1552,9 @@ GraphExecutor::CachedSegOpr GraphExecutor::CreateCachedSegOpr(size_t topo_start, OpNode& op_node = op_nodes_[nid]; if (op_node.skip_exec_node) continue; if (inode.source->is_variable()) continue; + // We shouldn't add control flow operators to a segment. + // We can't execute these operators in the engine. + if (op_node.exec->HasSubgraph()) continue; if (op_node.exec->exec_type() != ExecType::kSync) { return ret; } diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 70cb1ecddc2e..0d6c4039f62b 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5670,10 +5670,11 @@ def test_foreach(): v3 = mx.sym.var("v3") v4 = mx.sym.var("v4") g = v1 + v2 - op = mx.sym.Foreach(g, v3, v4) + out = mx.sym.Foreach(g, v3, v4) + out = out * 2 arr1 = mx.nd.random.uniform(shape=(5, 2)) arr2 = mx.nd.random.uniform(shape=(5, 2)) - e = op.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}) + e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}) e.forward() for y in e.outputs: y.wait_to_read() From 42ac88a8100364ea64db92eb57b16b475c8c791f Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Thu, 5 Apr 2018 23:05:51 +0000 Subject: [PATCH 07/71] Implements forward of foreach. --- src/operator/nn/control_flow.cc | 41 ++++++++++++++++++++++++-- tests/python/unittest/test_operator.py | 16 ++++++---- 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 614de336e938..5226290cd937 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -191,13 +191,45 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, g.attrs[ctx.is_train ? "full_mem_plan" : "forward_mem_plan"] = std::make_shared(std::move(mem_plan)); } - ExecSubgraph(g, ctx, inputs, req, outputs); + size_t len = inputs[0].shape()[0]; + CHECK_EQ(outputs.size(), 1U); + CHECK_EQ(inputs.size(), 2U); + CHECK_EQ(inputs[0].shape()[0], outputs[0].shape()[0]); + std::vector subg_inputs(inputs.size()); + std::vector subg_outputs(outputs.size()); + for (size_t i = 1; i < inputs.size(); i++) + subg_inputs[i] = inputs[i]; + // Here we iterate over the first dimension of the first input array. + for (size_t i = 0; i < len; i++) { + subg_inputs[0] = inputs[0].At(i); + // For the first iteration, the second argument is the second input array, + // i.e., the initial state. + if (i == 0) + subg_inputs[1] = inputs[1]; + else + // For the rest of the iterations, the second argument is the output from + // the previous iteration. + subg_inputs[1] = subg_outputs[0]; + subg_outputs[0] = outputs[0].At(i); + + ExecSubgraph(g, ctx, subg_inputs, req, subg_outputs); + // We need to wait for the iteration to complete before executing + // the next one or return from the loop. In this way, we can reuse + // the memory in the subgraph. + for (size_t j = 0; j < subg_outputs.size(); j++) + subg_outputs[j].WaitToRead(); + } } static bool ForeachShape(const nnvm::NodeAttrs& attrs, std::vector *in_shape, std::vector *out_shape) { + CHECK_EQ(in_shape->size(), 2U); nnvm::ShapeVector shape_inputs = *in_shape; + // foreach iterates over the first input NDArray over the first dimension. + shape_inputs[0] = TShape(in_shape->at(0).begin() + 1, in_shape->at(0).end()); + bool ret = shape_assign(&shape_inputs[1], shape_inputs[0]); + CHECK(ret); auto g = attrs.g; CHECK(g); // TODO(zhengda) This can also be called in the execution engine. @@ -206,7 +238,12 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, const auto& shapes = g->GetAttr("shape"); CHECK(g->outputs.size() == 1); uint32_t eid = g->indexed_graph().entry_id(g->outputs[0]); - (*out_shape)[0] = shapes[eid]; + const auto& g_out_shape = shapes[eid]; + const auto& in0 = (*in_shape)[0]; + CHECK_EQ(g_out_shape.ndim() + 1, in0.ndim()); + for (size_t i = 1; i < in0.ndim(); i++) + CHECK_EQ(in0[i], g_out_shape[i - 1]); + (*out_shape)[0] = in0; return true; } diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 0d6c4039f62b..a2d108824f99 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5673,13 +5673,19 @@ def test_foreach(): out = mx.sym.Foreach(g, v3, v4) out = out * 2 arr1 = mx.nd.random.uniform(shape=(5, 2)) - arr2 = mx.nd.random.uniform(shape=(5, 2)) + arr2 = mx.nd.random.uniform(shape=(2)) e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}) e.forward() - for y in e.outputs: - y.wait_to_read() - print(y) - print(arr1 + arr2) + arr1 = arr1.asnumpy() + arr2 = arr2.asnumpy() + np_res = np.zeros_like(arr1) + for i in range(arr1.shape[0]): + if (i == 0): + np_res[i] = arr2 + arr1[i] + else: + np_res[i] = np_res[i - 1] + arr1[i] + np_res = np_res * 2 + assert_almost_equal(e.outputs[0].asnumpy(), np_res, rtol=0.001, atol=0.0001) @with_seed() From 3318797949ff644e6501794f35a23f12370408ca Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 6 Apr 2018 01:12:22 +0000 Subject: [PATCH 08/71] Add support for variable numbers of inputs and outputs. --- src/operator/nn/control_flow.cc | 131 ++++++++++++++++++------- tests/python/unittest/test_operator.py | 5 +- 2 files changed, 99 insertions(+), 37 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 5226290cd937..0ff7201e03bc 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -158,6 +158,19 @@ static void ExecSubgraph(nnvm::Graph &g, const OpContext& ctx, std::move(ref_count), &states, dispatch_modes); } +struct ForeachParam : public dmlc::Parameter { + int num_args; + int dim; + DMLC_DECLARE_PARAMETER(ForeachParam) { + DMLC_DECLARE_FIELD(num_args).set_lower_bound(1) + .describe("Number of inputs."); + DMLC_DECLARE_FIELD(dim).set_default(1) + .describe("the dimension of the input array to iterate."); + } +}; // struct ForeachParam + +DMLC_REGISTER_PARAMETER(ForeachParam); + static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, const OpContext& ctx, const std::vector& inputs, @@ -192,58 +205,95 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, = std::make_shared(std::move(mem_plan)); } size_t len = inputs[0].shape()[0]; - CHECK_EQ(outputs.size(), 1U); - CHECK_EQ(inputs.size(), 2U); CHECK_EQ(inputs[0].shape()[0], outputs[0].shape()[0]); + + // Initialize the inputs for the subgraph. std::vector subg_inputs(inputs.size()); - std::vector subg_outputs(outputs.size()); - for (size_t i = 1; i < inputs.size(); i++) + for (size_t i = 1; i < inputs.size(); i++) { + // These are the initial states. subg_inputs[i] = inputs[i]; + } + + // Initialize the outputs of the subgraph is a little trickier. + // The states from the previous iteration are used as the inputs of the next + // iteration, so I have to maintain two arrays, so the inputs and outputs + // of the subgraph share the same memory. + std::vector subg_outputs1(inputs.size()); + std::vector subg_outputs2(inputs.size()); + std::vector *subg_outputs[2]{&subg_outputs1, &subg_outputs2}; + // If the length is an odd number, the last iteration will use the first set + // of outputs. In this way, we don't need to copy the results from the + // subgraph to the final outputs of the loop. + if (len % 2 == 1) { + for (size_t i = 1; i < subg_outputs1.size(); i++) { + subg_outputs1[i] = outputs[i]; + subg_outputs2[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), false, + outputs[i].dtype()); + } + } else { + // Otherwise, we'll use the second set of outputs. + for (size_t i = 1; i < subg_outputs1.size(); i++) { + subg_outputs1[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), false, + outputs[i].dtype()); + subg_outputs2[i] = outputs[i]; + } + } + // Here we iterate over the first dimension of the first input array. for (size_t i = 0; i < len; i++) { + std::vector *subg_out_curr = subg_outputs[i % 2]; + std::vector *subg_out_prev = subg_outputs[(i + 1) % 2]; + (*subg_out_curr)[0] = outputs[0].At(i); + + // Get a slice from the first input array. subg_inputs[0] = inputs[0].At(i); - // For the first iteration, the second argument is the second input array, - // i.e., the initial state. - if (i == 0) - subg_inputs[1] = inputs[1]; - else - // For the rest of the iterations, the second argument is the output from - // the previous iteration. - subg_inputs[1] = subg_outputs[0]; - subg_outputs[0] = outputs[0].At(i); + // For the rest of the iterations, the rest of the arguments are the outputs + // from the previous iteration. + if (i > 0) { + for (size_t j = 1; j < subg_out_prev->size(); j++) + subg_inputs[j] = (*subg_out_prev)[j]; + } - ExecSubgraph(g, ctx, subg_inputs, req, subg_outputs); + ExecSubgraph(g, ctx, subg_inputs, req, *subg_out_curr); // We need to wait for the iteration to complete before executing // the next one or return from the loop. In this way, we can reuse // the memory in the subgraph. - for (size_t j = 0; j < subg_outputs.size(); j++) - subg_outputs[j].WaitToRead(); + for (size_t j = 0; j < subg_out_curr->size(); j++) + (*subg_out_curr)[j].WaitToRead(); } } static bool ForeachShape(const nnvm::NodeAttrs& attrs, std::vector *in_shape, std::vector *out_shape) { - CHECK_EQ(in_shape->size(), 2U); nnvm::ShapeVector shape_inputs = *in_shape; // foreach iterates over the first input NDArray over the first dimension. shape_inputs[0] = TShape(in_shape->at(0).begin() + 1, in_shape->at(0).end()); - bool ret = shape_assign(&shape_inputs[1], shape_inputs[0]); - CHECK(ret); auto g = attrs.g; CHECK(g); + const auto& idx = g->indexed_graph(); + CHECK_EQ(idx.input_nodes().size(), in_shape->size()); + CHECK_EQ(idx.outputs().size(), out_shape->size()); // TODO(zhengda) This can also be called in the execution engine. // We need to make it thread-safe. imperative::CheckAndInferShape(g.get(), std::move(shape_inputs), true); const auto& shapes = g->GetAttr("shape"); - CHECK(g->outputs.size() == 1); - uint32_t eid = g->indexed_graph().entry_id(g->outputs[0]); + + // For the first shape. + uint32_t eid = idx.entry_id(g->outputs[0]); const auto& g_out_shape = shapes[eid]; - const auto& in0 = (*in_shape)[0]; + const auto &in0 = (*in_shape)[0]; + auto &out0 = (*out_shape)[0]; CHECK_EQ(g_out_shape.ndim() + 1, in0.ndim()); - for (size_t i = 1; i < in0.ndim(); i++) - CHECK_EQ(in0[i], g_out_shape[i - 1]); - (*out_shape)[0] = in0; + out0 = in0; + for (size_t i = 1; i < out0.ndim(); i++) + out0[i] = g_out_shape[i - 1]; + + // For the remaining shapes. + for (size_t i = 1; i < g->outputs.size(); i++) { + uint32_t eid = idx.entry_id(g->outputs[i]); + (*out_shape)[i] = shapes[eid]; + } return true; } @@ -252,13 +302,15 @@ static bool ForeachType(const nnvm::NodeAttrs& attrs, nnvm::DTypeVector dtype_inputs = *in_type; auto g = attrs.g; CHECK(g); + const auto& idx = g->indexed_graph(); + CHECK_EQ(idx.input_nodes().size(), in_type->size()); + CHECK_EQ(idx.outputs().size(), out_type->size()); // TODO(zhengda) This can also be called in the execution engine. // We need to make it thread-safe. imperative::CheckAndInferType(g.get(), std::move(dtype_inputs), true); const auto &dtypes = g->GetAttr("dtype"); - CHECK(g->outputs.size() == 1); - uint32_t eid = g->indexed_graph().entry_id(g->outputs[0]); - (*out_type)[0] = dtypes[eid]; + for (size_t i = 0; i < g->outputs.size(); i++) + (*out_type)[i] = dtypes[idx.entry_id(g->outputs[i])]; return true; } @@ -270,7 +322,8 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, auto g = attrs.g; CHECK(g); const auto& idx = g->indexed_graph(); - CHECK(idx.input_nodes().size() == in_attrs->size()); + CHECK_EQ(idx.input_nodes().size(), in_attrs->size()); + CHECK_EQ(idx.outputs().size(), out_attrs->size()); exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); StorageTypeVector storage_type_inputs = *in_attrs; imperative::CheckAndInferStorageType(g.get(), std::move(dev_masks), @@ -279,18 +332,23 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, const auto& stypes = g->GetAttr("storage_type"); auto &outputs = idx.outputs(); CHECK(outputs.size() == out_attrs->size()); - for (size_t i = 0; i < out_attrs->size(); i++) { + for (size_t i = 0; i < out_attrs->size(); i++) (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; - } return true; } NNVM_REGISTER_OP(Foreach) .describe(R"code(Foreach)code" ADD_FILELINE) -//.set_attr_parser(ParamParser) +.set_attr_parser(ParamParser) .set_attr("FInferStorageType", ForeachStorageType) -.set_num_inputs(3) -.set_num_outputs(1) +.set_num_inputs([](const NodeAttrs& attrs) { + const ForeachParam& params = nnvm::get(attrs.parsed); + return params.num_args; +}) +.set_num_outputs([](const NodeAttrs& attrs) { + const ForeachParam& params = nnvm::get(attrs.parsed); + return params.num_args - 1; +}) .set_attr("FListInputNames", [](const NodeAttrs& attrs) { return std::vector{"fn", "data1", "data2"}; @@ -302,9 +360,10 @@ NNVM_REGISTER_OP(Foreach) .set_attr("FInferShape", ForeachShape) .set_attr("FInferType", ForeachType) .set_attr("FComputeEx", ForeachComputeExCPU) +.set_attr("key_var_num_args", "num_args") .add_argument("fn", "Symbol", "Input graph.") -.add_argument("data1", "NDArray-or-Symbol", "Input1.") -.add_argument("data2", "NDArray-or-Symbol", "Input2."); +.add_argument("input", "NDArray-or-Symbol", "The input array where we iterate over.") +.add_argument("states", "NDArray-or-Symbol[]", "The list of initial states."); //.add_arguments(ForeachParam::__FIELDS__()); } // namespace op diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index a2d108824f99..fe6e6a6d0ee8 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5670,8 +5670,11 @@ def test_foreach(): v3 = mx.sym.var("v3") v4 = mx.sym.var("v4") g = v1 + v2 + # TODO This is problematic. We can't count on the user to define two different symbols. + g = mx.sym.Group([g, g * 1]) out = mx.sym.Foreach(g, v3, v4) - out = out * 2 + out1 = out[0] * 2 + out = mx.sym.Group([out1, out[1]]) arr1 = mx.nd.random.uniform(shape=(5, 2)) arr2 = mx.nd.random.uniform(shape=(2)) e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}) From 7670aceea6b2ffdb2e19ff57c819165ed6a850b4 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 6 Apr 2018 19:05:30 +0000 Subject: [PATCH 09/71] Add a python wrapper for foreach. --- python/mxnet/contrib/__init__.py | 1 + python/mxnet/contrib/control_flow.py | 39 ++++++++++++++++++++++++++ tests/python/unittest/test_operator.py | 13 +++++---- 3 files changed, 47 insertions(+), 6 deletions(-) create mode 100644 python/mxnet/contrib/control_flow.py diff --git a/python/mxnet/contrib/__init__.py b/python/mxnet/contrib/__init__.py index fbfd3469678b..7489d97d90fe 100644 --- a/python/mxnet/contrib/__init__.py +++ b/python/mxnet/contrib/__init__.py @@ -32,3 +32,4 @@ from . import io from . import quantization from . import quantization as quant +from . import control_flow as cf diff --git a/python/mxnet/contrib/control_flow.py b/python/mxnet/contrib/control_flow.py new file mode 100644 index 000000000000..8b128fbb8d77 --- /dev/null +++ b/python/mxnet/contrib/control_flow.py @@ -0,0 +1,39 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import mxnet as mx + +def foreach(func, input, init_states, back_prop=False): + in_ele = mx.sym.var("in") + states = [] + i = 0 + assert isinstance(init_states, list), "init_states should be a list" + for s in init_states: + states.append(mx.sym.var("state" + str(i))) + i = i + 1 + sym_out = func(in_ele, states) + # The function should return a tuple. The first element goes to + # the output of the function. The second element is a list. + assert isinstance(sym_out, tuple), "func should return a tuple (out, states)" + assert isinstance(sym_out[1], list), \ + "the second element in the returned tuple should be a list" + + flat_out = [sym_out[0]] + for s in sym_out[1]: + flat_out.append(s) + g = mx.sym.Group(flat_out) + return mx.sym.contrib.Foreach(g, input, *init_states) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index fe6e6a6d0ee8..224ad8b0e926 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5665,14 +5665,15 @@ def test_float16_min_max(): @with_seed() def test_foreach(): - v1 = mx.sym.var("v1") - v2 = mx.sym.var("v2") v3 = mx.sym.var("v3") v4 = mx.sym.var("v4") - g = v1 + v2 - # TODO This is problematic. We can't count on the user to define two different symbols. - g = mx.sym.Group([g, g * 1]) - out = mx.sym.Foreach(g, v3, v4) + + def step(in1, states): + out = in1 + states[0] + # TODO This is problematic. We can't count on the user to define two different symbols. + return (out, [out * 1]) + + out = mx.contrib.cf.foreach(step, v3, [v4]) out1 = out[0] * 2 out = mx.sym.Group([out1, out[1]]) arr1 = mx.nd.random.uniform(shape=(5, 2)) From d98f87836e0137d9d79171054ecb7c2c71abf2ad Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 6 Apr 2018 23:57:35 +0000 Subject: [PATCH 10/71] Fix the order of inputs. --- src/operator/nn/control_flow.cc | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 0ff7201e03bc..7d59fa7d5c30 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -171,6 +171,26 @@ struct ForeachParam : public dmlc::Parameter { DMLC_REGISTER_PARAMETER(ForeachParam); +// The input arguments are ordered in the following order: +// in, state0, state1, ... +// We need to reorder them in the same order as the input nodes of the subgraph. +template +static std::vector ReorderInputs(const std::vector &in, const nnvm::IndexedGraph& idx) { + std::vector ret(in.size()); + CHECK_EQ(idx.input_nodes().size(), in.size()); + for (size_t i = 0; i < idx.input_nodes().size(); i++) { + std::string name = idx[idx.input_nodes()[i]].source->attrs.name; + if (name == "in") { + ret[i] = in[0]; + } else { + auto idx_str = name.substr(5); + int idx = std::stoi(idx_str); + ret[i] = in[idx + 1]; + } + } + return ret; +} + static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, const OpContext& ctx, const std::vector& inputs, @@ -178,6 +198,7 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, const std::vector& outputs) { CHECK(attrs.g != nullptr); nnvm::Graph &g = *attrs.g; + const auto& idx = g.indexed_graph(); // If this is inference, we only need the forward memory plan. bool has_mem_plan = !ctx.is_train && g.attrs.count("forward_mem_plan"); @@ -185,7 +206,6 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, has_mem_plan = has_mem_plan || (ctx.is_train && g.attrs.count("full_mem_plan")); // If we don't have a memory plan yet, we need to create a memory plan. if (!has_mem_plan) { - const auto& idx = g.indexed_graph(); nnvm::StorageVector storage(idx.num_node_entries(), exec::kBadStorageID); for (const auto i : idx.input_nodes()) storage[idx.entry_id(i, 0)] = exec::kExternalStorageID; @@ -254,7 +274,8 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, subg_inputs[j] = (*subg_out_prev)[j]; } - ExecSubgraph(g, ctx, subg_inputs, req, *subg_out_curr); + std::vector reordered_ins = ReorderInputs(subg_inputs, idx); + ExecSubgraph(g, ctx, reordered_ins, req, *subg_out_curr); // We need to wait for the iteration to complete before executing // the next one or return from the loop. In this way, we can reuse // the memory in the subgraph. @@ -276,6 +297,7 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, CHECK_EQ(idx.outputs().size(), out_shape->size()); // TODO(zhengda) This can also be called in the execution engine. // We need to make it thread-safe. + shape_inputs = ReorderInputs(shape_inputs, idx); imperative::CheckAndInferShape(g.get(), std::move(shape_inputs), true); const auto& shapes = g->GetAttr("shape"); @@ -307,6 +329,7 @@ static bool ForeachType(const nnvm::NodeAttrs& attrs, CHECK_EQ(idx.outputs().size(), out_type->size()); // TODO(zhengda) This can also be called in the execution engine. // We need to make it thread-safe. + dtype_inputs = ReorderInputs(dtype_inputs, idx); imperative::CheckAndInferType(g.get(), std::move(dtype_inputs), true); const auto &dtypes = g->GetAttr("dtype"); for (size_t i = 0; i < g->outputs.size(); i++) @@ -326,6 +349,7 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, CHECK_EQ(idx.outputs().size(), out_attrs->size()); exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); StorageTypeVector storage_type_inputs = *in_attrs; + storage_type_inputs = ReorderInputs(storage_type_inputs, idx); imperative::CheckAndInferStorageType(g.get(), std::move(dev_masks), std::move(storage_type_inputs), true); *dispatch_mode = DispatchMode::kFComputeEx; From 4d6779c16316d42d97e12440e5fc101253bd6492 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 7 Apr 2018 00:06:49 +0000 Subject: [PATCH 11/71] hide C version of foreach. --- python/mxnet/contrib/control_flow.py | 2 +- src/operator/nn/control_flow.cc | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/mxnet/contrib/control_flow.py b/python/mxnet/contrib/control_flow.py index 8b128fbb8d77..b6705ca4552a 100644 --- a/python/mxnet/contrib/control_flow.py +++ b/python/mxnet/contrib/control_flow.py @@ -36,4 +36,4 @@ def foreach(func, input, init_states, back_prop=False): for s in sym_out[1]: flat_out.append(s) g = mx.sym.Group(flat_out) - return mx.sym.contrib.Foreach(g, input, *init_states) + return mx.sym._internal._foreach(g, input, *init_states) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 7d59fa7d5c30..e95b88df2ad8 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -361,8 +361,8 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, return true; } -NNVM_REGISTER_OP(Foreach) -.describe(R"code(Foreach)code" ADD_FILELINE) +NNVM_REGISTER_OP(_foreach) +.describe(R"code(foreach)code" ADD_FILELINE) .set_attr_parser(ParamParser) .set_attr("FInferStorageType", ForeachStorageType) .set_num_inputs([](const NodeAttrs& attrs) { From 43f8c1b8de0c9e4928fa8dc349656d2ee124ca8f Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 7 Apr 2018 01:11:26 +0000 Subject: [PATCH 12/71] fix a bug temporarily. --- python/mxnet/contrib/control_flow.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/mxnet/contrib/control_flow.py b/python/mxnet/contrib/control_flow.py index b6705ca4552a..42d869208898 100644 --- a/python/mxnet/contrib/control_flow.py +++ b/python/mxnet/contrib/control_flow.py @@ -34,6 +34,9 @@ def foreach(func, input, init_states, back_prop=False): flat_out = [sym_out[0]] for s in sym_out[1]: - flat_out.append(s) + # There is a problem if the outputs are the same as the inputs + # or the first output. + # TODO this is a temp fix. + flat_out.append(mx.sym.identity(s)) g = mx.sym.Group(flat_out) return mx.sym._internal._foreach(g, input, *init_states) From efb6ba0d26d601de360ed43c90383c900910336f Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 6 Apr 2018 23:57:57 +0000 Subject: [PATCH 13/71] add test with lstm. --- tests/python/unittest/test_operator.py | 57 ++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 224ad8b0e926..4781f92c7f76 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5668,6 +5668,7 @@ def test_foreach(): v3 = mx.sym.var("v3") v4 = mx.sym.var("v4") + # This tests foreach with accumulation sum. def step(in1, states): out = in1 + states[0] # TODO This is problematic. We can't count on the user to define two different symbols. @@ -5692,6 +5693,62 @@ def step(in1, states): assert_almost_equal(e.outputs[0].asnumpy(), np_res, rtol=0.001, atol=0.0001) +@with_seed() +def test_foreach_lstm(): + # This tests foreach with accumulation sum. + def step(in1, states): + params = mx.rnn.RNNParams() + params._params['i2h_weight'] = states[2] + params._params['h2h_weight'] = states[3] + params._params['i2h_bias'] = states[4] + params._params['h2h_bias'] = states[5] + lstm = mx.rnn.LSTMCell(4, prefix='mylstm_', params=params) + prev_states = [states[0], states[1]] + next_h, [next_h, next_c] = lstm(in1, prev_states) + # TODO This is problematic. We can't count on the user to define two different symbols. + return (next_h, [next_h, next_c, states[2], states[3], states[4], states[5]]) + + data = mx.sym.var("data") + init_h = mx.sym.var("h") + init_c = mx.sym.var("c") + i2h_weight = mx.sym.var("i2h_weight") + h2h_weight = mx.sym.var("h2h_weight") + i2h_bias = mx.sym.var("i2h_bias") + h2h_bias = mx.sym.var("h2h_bias") + + data_arr = mx.nd.random.uniform(shape=(5, 2, 4)) + h_arr = mx.nd.random.uniform(shape=(2, 4)) + c_arr = mx.nd.random.uniform(shape=(2, 4)) + i2h_warr = mx.nd.random.uniform(shape=(16, 4)) + h2h_warr = mx.nd.random.uniform(shape=(16, 4)) + i2h_barr = mx.nd.random.uniform(shape=(16)) + h2h_barr = mx.nd.random.uniform(shape=(16)) + + out = mx.contrib.cf.foreach(step, data, [init_h, init_c, i2h_weight, h2h_weight, i2h_bias, h2h_bias]) + e = out.bind(ctx=mx.cpu(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, + 'i2h_weight': i2h_warr, 'h2h_weight': h2h_warr, 'i2h_bias': i2h_barr, 'h2h_bias': h2h_barr}) + e.forward() + outputs1 = e.outputs + + lstm = mx.rnn.LSTMCell(4, prefix='mylstm_') + h = init_h + c = init_c + unroll_outs = [] + for inputs in mx.sym.split(data, num_outputs=data_arr.shape[0], axis=0, squeeze_axis=True): + h, [h, c] = lstm(inputs, [h, c]) + unroll_outs.append(mx.sym.expand_dims(h, axis=0)) + unroll_outs = mx.sym.concat(*unroll_outs, dim=0) + out = mx.sym.Group([unroll_outs, h, c]) + e = out.bind(ctx=mx.cpu(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, + 'mylstm_i2h_weight': i2h_warr, 'mylstm_h2h_weight': h2h_warr, + 'mylstm_i2h_bias': i2h_barr, 'mylstm_h2h_bias': h2h_barr}) + e.forward() + outputs2 = e.outputs + + for i in range(len(outputs2)): + assert_almost_equal(outputs1[i].asnumpy(), outputs2[i].asnumpy(), rtol=0.001, atol=0.0001) + + @with_seed() def test_squeeze_op(): def check_squeeze_op(shape, axis=None): From 7593275e74a0e987705e2ecbed379be293d744ad Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 7 Apr 2018 01:45:37 +0000 Subject: [PATCH 14/71] Test free variables. --- python/mxnet/contrib/control_flow.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/mxnet/contrib/control_flow.py b/python/mxnet/contrib/control_flow.py index 42d869208898..781d6c7be9b1 100644 --- a/python/mxnet/contrib/control_flow.py +++ b/python/mxnet/contrib/control_flow.py @@ -19,11 +19,13 @@ def foreach(func, input, init_states, back_prop=False): in_ele = mx.sym.var("in") + gin_names = ["in"] states = [] i = 0 assert isinstance(init_states, list), "init_states should be a list" for s in init_states: states.append(mx.sym.var("state" + str(i))) + gin_names.append("state" + str(i)) i = i + 1 sym_out = func(in_ele, states) # The function should return a tuple. The first element goes to @@ -39,4 +41,9 @@ def foreach(func, input, init_states, back_prop=False): # TODO this is a temp fix. flat_out.append(mx.sym.identity(s)) g = mx.sym.Group(flat_out) + + # The input function can't have free variables right now. + for i in g.list_inputs(): + assert i in gin_names, "The input function can't contain free variables" + return mx.sym._internal._foreach(g, input, *init_states) From 6d4b90bb101f29e895c59c0f0a30aa5c2404f083 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 9 Apr 2018 13:09:56 -0700 Subject: [PATCH 15/71] change for the new interface of InputGraph attribute. --- src/executor/attach_op_execs_pass.cc | 4 ++-- src/operator/nn/control_flow.cc | 15 +++++++++------ 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/executor/attach_op_execs_pass.cc b/src/executor/attach_op_execs_pass.cc index b78149975113..1944927d66a9 100644 --- a/src/executor/attach_op_execs_pass.cc +++ b/src/executor/attach_op_execs_pass.cc @@ -202,7 +202,7 @@ class FComputeExecutor : public StorageFallbackOpExecutor { } bool HasSubgraph() const override { - return attrs_.g != nullptr; + return !attrs_.subgraphs.empty(); } explicit FComputeExecutor(const NodeAttrs& attrs, FCompute fcompute, @@ -231,7 +231,7 @@ class FComputeExExecutor : public OpExecutor { void Setup() override {} bool HasSubgraph() const override { - return attrs_.g != nullptr; + return !attrs_.subgraphs.empty(); } ExecType exec_type() const override { diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index e95b88df2ad8..99a12a9fa901 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -196,8 +196,8 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, const std::vector& inputs, const std::vector& req, const std::vector& outputs) { - CHECK(attrs.g != nullptr); - nnvm::Graph &g = *attrs.g; + CHECK_EQ(attrs.subgraphs.size(), 1U); + nnvm::Graph &g = *attrs.subgraphs[0]; const auto& idx = g.indexed_graph(); // If this is inference, we only need the forward memory plan. @@ -290,7 +290,8 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, nnvm::ShapeVector shape_inputs = *in_shape; // foreach iterates over the first input NDArray over the first dimension. shape_inputs[0] = TShape(in_shape->at(0).begin() + 1, in_shape->at(0).end()); - auto g = attrs.g; + CHECK_EQ(attrs.subgraphs.size(), 1U); + auto g = attrs.subgraphs[0]; CHECK(g); const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_shape->size()); @@ -322,7 +323,8 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, static bool ForeachType(const nnvm::NodeAttrs& attrs, std::vector *in_type, std::vector *out_type) { nnvm::DTypeVector dtype_inputs = *in_type; - auto g = attrs.g; + CHECK_EQ(attrs.subgraphs.size(), 1U); + auto g = attrs.subgraphs[0]; CHECK(g); const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_type->size()); @@ -342,7 +344,8 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, DispatchMode* dispatch_mode, std::vector *in_attrs, std::vector *out_attrs) { - auto g = attrs.g; + CHECK_EQ(attrs.subgraphs.size(), 1U); + auto g = attrs.subgraphs[0]; CHECK(g); const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_attrs->size()); @@ -379,7 +382,7 @@ NNVM_REGISTER_OP(_foreach) }) .set_attr("FInputGraph", [](const NodeAttrs& attrs) { - return 0; + return std::vector{0}; }) .set_attr("FInferShape", ForeachShape) .set_attr("FInferType", ForeachType) From bde96cd3f45754396db49810067a9220ac669919 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 11 Apr 2018 21:04:40 +0000 Subject: [PATCH 16/71] Add attribute to the subgraph. --- python/mxnet/contrib/control_flow.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/python/mxnet/contrib/control_flow.py b/python/mxnet/contrib/control_flow.py index 781d6c7be9b1..a839c467ad7e 100644 --- a/python/mxnet/contrib/control_flow.py +++ b/python/mxnet/contrib/control_flow.py @@ -17,7 +17,7 @@ import mxnet as mx -def foreach(func, input, init_states, back_prop=False): +def foreach(func, input, init_states, back_prop=False, name="foreach"): in_ele = mx.sym.var("in") gin_names = ["in"] states = [] @@ -27,19 +27,20 @@ def foreach(func, input, init_states, back_prop=False): states.append(mx.sym.var("state" + str(i))) gin_names.append("state" + str(i)) i = i + 1 - sym_out = func(in_ele, states) - # The function should return a tuple. The first element goes to - # the output of the function. The second element is a list. - assert isinstance(sym_out, tuple), "func should return a tuple (out, states)" - assert isinstance(sym_out[1], list), \ - "the second element in the returned tuple should be a list" + with mx.AttrScope(subgraph_name=name): + sym_out = func(in_ele, states) + # The function should return a tuple. The first element goes to + # the output of the function. The second element is a list. + assert isinstance(sym_out, tuple), "func should return a tuple (out, states)" + assert isinstance(sym_out[1], list), \ + "the second element in the returned tuple should be a list" - flat_out = [sym_out[0]] - for s in sym_out[1]: - # There is a problem if the outputs are the same as the inputs - # or the first output. - # TODO this is a temp fix. - flat_out.append(mx.sym.identity(s)) + flat_out = [sym_out[0]] + for s in sym_out[1]: + # There is a problem if the outputs are the same as the inputs + # or the first output. + # TODO this is a temp fix. + flat_out.append(mx.sym.identity(s)) g = mx.sym.Group(flat_out) # The input function can't have free variables right now. From 8bb020e8d9a46de3e4a69d42ab306d74ea83736d Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 11 Apr 2018 22:21:20 +0000 Subject: [PATCH 17/71] Handle free variables. --- python/mxnet/contrib/control_flow.py | 39 ++++++++++++++++++++++---- src/c_api/c_api_symbolic.cc | 5 ++-- src/operator/nn/control_flow.cc | 36 ++++++++++++++++-------- tests/python/unittest/test_operator.py | 27 +++++++++--------- 4 files changed, 73 insertions(+), 34 deletions(-) diff --git a/python/mxnet/contrib/control_flow.py b/python/mxnet/contrib/control_flow.py index a839c467ad7e..2e3a8997b016 100644 --- a/python/mxnet/contrib/control_flow.py +++ b/python/mxnet/contrib/control_flow.py @@ -24,8 +24,8 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): i = 0 assert isinstance(init_states, list), "init_states should be a list" for s in init_states: - states.append(mx.sym.var("state" + str(i))) - gin_names.append("state" + str(i)) + states.append(mx.sym.var(s.name)) + gin_names.append(s.name) i = i + 1 with mx.AttrScope(subgraph_name=name): sym_out = func(in_ele, states) @@ -43,8 +43,35 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): flat_out.append(mx.sym.identity(s)) g = mx.sym.Group(flat_out) - # The input function can't have free variables right now. - for i in g.list_inputs(): - assert i in gin_names, "The input function can't contain free variables" + # Find free variables in the python that are symbols. + freevars = dict(zip(func.func_code.co_freevars, + (c.cell_contents for c in func.func_closure))) + sym_freevars = [] + for name in freevars: + val = freevars[name] + if isinstance(val, mx.sym.Symbol): + # We need to save the original symbol first. + sym_freevars.append(val) + gin_names.append(name) - return mx.sym._internal._foreach(g, input, *init_states) + if (isinstance(input, list)): + num_inputs = len(input) + else: + num_inputs = 1 + + # Here we need to find out how the input symbols are ordered as well as + # where the loop states are located in the list of inputs. + ins = init_states + sym_freevars + ins = {sym.name:sym for sym in ins} + ordered_ins = [] + in_state_locs = [-1] * len(init_states) + for in_name in g.list_inputs(): + assert in_name in gin_names, "The input graph contains variables we can't find" + if in_name in ins: + ordered_ins.append(ins[in_name]) + for i in range(len(init_states)): + if (init_states[i].name == in_name): + in_state_locs[i] = len(ordered_ins) - 1 + num_inputs + + return mx.sym._internal._foreach(g, input, *ordered_ins, num_outputs=len(flat_out), + in_state_locs=in_state_locs) diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index 4666b6adf0c3..6ba8a3d95469 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -38,10 +38,11 @@ void RegisterLegacyOpProp(); void RegisterLegacyNDFunc(); } const std::vector kHiddenKeys = { - "ctx_group", "lr_mult", "wd_mult", "force_mirroring", "mirror_stage" + "ctx_group", "lr_mult", "wd_mult", "force_mirroring", "mirror_stage", "subgraph_name" }; const std::vector kReplacedHiddenKeys = { - "__ctx_group__", "__lr_mult__", "__wd_mult__", "__force_mirroring__", "__mirror_stage__" + "__ctx_group__", "__lr_mult__", "__wd_mult__", "__force_mirroring__", "__mirror_stage__", + "subgraph_name" }; const char *kNamespaceSeparator = "$"; diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 99a12a9fa901..135694742bed 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -161,11 +161,17 @@ static void ExecSubgraph(nnvm::Graph &g, const OpContext& ctx, struct ForeachParam : public dmlc::Parameter { int num_args; int dim; + int num_outputs; + nnvm::Tuple in_state_locs; DMLC_DECLARE_PARAMETER(ForeachParam) { DMLC_DECLARE_FIELD(num_args).set_lower_bound(1) .describe("Number of inputs."); DMLC_DECLARE_FIELD(dim).set_default(1) .describe("the dimension of the input array to iterate."); + DMLC_DECLARE_FIELD(num_outputs) + .describe("The number of outputs of the subgraph."); + DMLC_DECLARE_FIELD(in_state_locs) + .describe("The locations of loop states among the inputs."); } }; // struct ForeachParam @@ -196,6 +202,8 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, const std::vector& inputs, const std::vector& req, const std::vector& outputs) { + const ForeachParam& params = nnvm::get(attrs.parsed); + CHECK_EQ(outputs.size(), (size_t) params.num_outputs); CHECK_EQ(attrs.subgraphs.size(), 1U); nnvm::Graph &g = *attrs.subgraphs[0]; const auto& idx = g.indexed_graph(); @@ -238,8 +246,8 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, // The states from the previous iteration are used as the inputs of the next // iteration, so I have to maintain two arrays, so the inputs and outputs // of the subgraph share the same memory. - std::vector subg_outputs1(inputs.size()); - std::vector subg_outputs2(inputs.size()); + std::vector subg_outputs1(outputs.size()); + std::vector subg_outputs2(outputs.size()); std::vector *subg_outputs[2]{&subg_outputs1, &subg_outputs2}; // If the length is an odd number, the last iteration will use the first set // of outputs. In this way, we don't need to copy the results from the @@ -270,12 +278,13 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, // For the rest of the iterations, the rest of the arguments are the outputs // from the previous iteration. if (i > 0) { - for (size_t j = 1; j < subg_out_prev->size(); j++) - subg_inputs[j] = (*subg_out_prev)[j]; + for (size_t j = 1; j < subg_out_prev->size(); j++) { + CHECK_LT(params.in_state_locs[j - 1], subg_inputs.size()); + subg_inputs[params.in_state_locs[j - 1]] = (*subg_out_prev)[j]; + } } - std::vector reordered_ins = ReorderInputs(subg_inputs, idx); - ExecSubgraph(g, ctx, reordered_ins, req, *subg_out_curr); + ExecSubgraph(g, ctx, subg_inputs, req, *subg_out_curr); // We need to wait for the iteration to complete before executing // the next one or return from the loop. In this way, we can reuse // the memory in the subgraph. @@ -287,6 +296,8 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, static bool ForeachShape(const nnvm::NodeAttrs& attrs, std::vector *in_shape, std::vector *out_shape) { + const ForeachParam& params = nnvm::get(attrs.parsed); + CHECK_EQ(out_shape->size(), (size_t) params.num_outputs); nnvm::ShapeVector shape_inputs = *in_shape; // foreach iterates over the first input NDArray over the first dimension. shape_inputs[0] = TShape(in_shape->at(0).begin() + 1, in_shape->at(0).end()); @@ -298,7 +309,6 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, CHECK_EQ(idx.outputs().size(), out_shape->size()); // TODO(zhengda) This can also be called in the execution engine. // We need to make it thread-safe. - shape_inputs = ReorderInputs(shape_inputs, idx); imperative::CheckAndInferShape(g.get(), std::move(shape_inputs), true); const auto& shapes = g->GetAttr("shape"); @@ -322,6 +332,8 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, static bool ForeachType(const nnvm::NodeAttrs& attrs, std::vector *in_type, std::vector *out_type) { + const ForeachParam& params = nnvm::get(attrs.parsed); + CHECK_EQ(out_type->size(), (size_t) params.num_outputs); nnvm::DTypeVector dtype_inputs = *in_type; CHECK_EQ(attrs.subgraphs.size(), 1U); auto g = attrs.subgraphs[0]; @@ -331,7 +343,6 @@ static bool ForeachType(const nnvm::NodeAttrs& attrs, CHECK_EQ(idx.outputs().size(), out_type->size()); // TODO(zhengda) This can also be called in the execution engine. // We need to make it thread-safe. - dtype_inputs = ReorderInputs(dtype_inputs, idx); imperative::CheckAndInferType(g.get(), std::move(dtype_inputs), true); const auto &dtypes = g->GetAttr("dtype"); for (size_t i = 0; i < g->outputs.size(); i++) @@ -344,6 +355,8 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, DispatchMode* dispatch_mode, std::vector *in_attrs, std::vector *out_attrs) { + const ForeachParam& params = nnvm::get(attrs.parsed); + CHECK_EQ(out_attrs->size(), (size_t) params.num_outputs); CHECK_EQ(attrs.subgraphs.size(), 1U); auto g = attrs.subgraphs[0]; CHECK(g); @@ -352,7 +365,6 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, CHECK_EQ(idx.outputs().size(), out_attrs->size()); exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); StorageTypeVector storage_type_inputs = *in_attrs; - storage_type_inputs = ReorderInputs(storage_type_inputs, idx); imperative::CheckAndInferStorageType(g.get(), std::move(dev_masks), std::move(storage_type_inputs), true); *dispatch_mode = DispatchMode::kFComputeEx; @@ -374,7 +386,7 @@ NNVM_REGISTER_OP(_foreach) }) .set_num_outputs([](const NodeAttrs& attrs) { const ForeachParam& params = nnvm::get(attrs.parsed); - return params.num_args - 1; + return params.num_outputs; }) .set_attr("FListInputNames", [](const NodeAttrs& attrs) { @@ -390,8 +402,8 @@ NNVM_REGISTER_OP(_foreach) .set_attr("key_var_num_args", "num_args") .add_argument("fn", "Symbol", "Input graph.") .add_argument("input", "NDArray-or-Symbol", "The input array where we iterate over.") -.add_argument("states", "NDArray-or-Symbol[]", "The list of initial states."); -//.add_arguments(ForeachParam::__FIELDS__()); +.add_argument("states", "NDArray-or-Symbol[]", "The list of initial states.") +.add_arguments(ForeachParam::__FIELDS__()); } // namespace op } // namespace mxnet diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 4781f92c7f76..2904c9c80ded 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5695,19 +5695,6 @@ def step(in1, states): @with_seed() def test_foreach_lstm(): - # This tests foreach with accumulation sum. - def step(in1, states): - params = mx.rnn.RNNParams() - params._params['i2h_weight'] = states[2] - params._params['h2h_weight'] = states[3] - params._params['i2h_bias'] = states[4] - params._params['h2h_bias'] = states[5] - lstm = mx.rnn.LSTMCell(4, prefix='mylstm_', params=params) - prev_states = [states[0], states[1]] - next_h, [next_h, next_c] = lstm(in1, prev_states) - # TODO This is problematic. We can't count on the user to define two different symbols. - return (next_h, [next_h, next_c, states[2], states[3], states[4], states[5]]) - data = mx.sym.var("data") init_h = mx.sym.var("h") init_c = mx.sym.var("c") @@ -5716,6 +5703,18 @@ def step(in1, states): i2h_bias = mx.sym.var("i2h_bias") h2h_bias = mx.sym.var("h2h_bias") + # This tests foreach with accumulation sum. + def step(in1, states): + params = mx.rnn.RNNParams() + params._params['i2h_weight'] = i2h_weight + params._params['h2h_weight'] = h2h_weight + params._params['i2h_bias'] = i2h_bias + params._params['h2h_bias'] = h2h_bias + lstm = mx.rnn.LSTMCell(4, prefix='mylstm_', params=params) + next_h, [next_h, next_c] = lstm(in1, states) + # TODO This is problematic. We can't count on the user to define two different symbols. + return (next_h, [next_h, next_c]) + data_arr = mx.nd.random.uniform(shape=(5, 2, 4)) h_arr = mx.nd.random.uniform(shape=(2, 4)) c_arr = mx.nd.random.uniform(shape=(2, 4)) @@ -5724,7 +5723,7 @@ def step(in1, states): i2h_barr = mx.nd.random.uniform(shape=(16)) h2h_barr = mx.nd.random.uniform(shape=(16)) - out = mx.contrib.cf.foreach(step, data, [init_h, init_c, i2h_weight, h2h_weight, i2h_bias, h2h_bias]) + out = mx.contrib.cf.foreach(step, data, [init_h, init_c]) e = out.bind(ctx=mx.cpu(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, 'i2h_weight': i2h_warr, 'h2h_weight': h2h_warr, 'i2h_bias': i2h_barr, 'h2h_bias': h2h_barr}) e.forward() From a4998f6a6b1cf6226047326e30ddf823a1d878a1 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 13 Apr 2018 18:43:14 +0000 Subject: [PATCH 18/71] Get all input symbols of a subgraph. --- include/mxnet/c_api.h | 10 ++++ python/mxnet/contrib/control_flow.py | 85 +++++++++++++++++----------- src/c_api/c_api_symbolic.cc | 27 +++++++++ 3 files changed, 89 insertions(+), 33 deletions(-) diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h index 06e39bfeb38b..4c37ba9d5400 100644 --- a/include/mxnet/c_api.h +++ b/include/mxnet/c_api.h @@ -1056,6 +1056,16 @@ MXNET_DLL int MXSymbolListAtomicSymbolCreators(mx_uint *out_size, */ MXNET_DLL int MXSymbolGetAtomicSymbolName(AtomicSymbolCreator creator, const char **name); + +/*! + * \brief Get the input symbols of the graph. + * \param sym The graph. + * \param outs The input symbols of the graph. + * \param out_size the number of input symbols returned. + */ +MXNET_DLL int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **outs, + int *out_size); + /*! * \brief Get the detailed information about atomic symbol. * \param creator the AtomicSymbolCreator. diff --git a/python/mxnet/contrib/control_flow.py b/python/mxnet/contrib/control_flow.py index 2e3a8997b016..6175df124eb1 100644 --- a/python/mxnet/contrib/control_flow.py +++ b/python/mxnet/contrib/control_flow.py @@ -15,44 +15,54 @@ # specific language governing permissions and limitations # under the License. -import mxnet as mx +import ctypes + +from .. import symbol +from ..base import _LIB, c_str, c_array, check_call +from ..base import SymbolHandle, NDArrayHandle +from ..attribute import AttrScope + +def _get_graph_inputs(subg, name, prefix): + num_handles = ctypes.c_int(1000) + handles = c_array(SymbolHandle, [SymbolHandle(0) for i in range(1000)]) + check_call(_LIB.MXSymbolGetInputSymbols(subg.handle, handles, + ctypes.byref(num_handles))) + + syms = [] + for i in range(num_handles.value): + s = symbol.Symbol(handles[i]) + syms.append(s) + return syms def foreach(func, input, init_states, back_prop=False, name="foreach"): - in_ele = mx.sym.var("in") - gin_names = ["in"] - states = [] - i = 0 assert isinstance(init_states, list), "init_states should be a list" - for s in init_states: - states.append(mx.sym.var(s.name)) - gin_names.append(s.name) - i = i + 1 - with mx.AttrScope(subgraph_name=name): + states = [] + with AttrScope(subgraph_name=name): + in_ele = symbol.var("in") + for s in init_states: + states.append(symbol.var(s.name)) + sym_out = func(in_ele, states) # The function should return a tuple. The first element goes to # the output of the function. The second element is a list. assert isinstance(sym_out, tuple), "func should return a tuple (out, states)" assert isinstance(sym_out[1], list), \ "the second element in the returned tuple should be a list" + assert len(sym_out[1]) == len(init_states), \ + "the number of output states (%d) should be the same as input states (%d)" \ + % (len(sym_out[1]), len(init_states)) - flat_out = [sym_out[0]] + if (isinstance(sym_out[0], list)): + flat_out = sym_out[0] + else: + flat_out = [sym_out[0]] for s in sym_out[1]: # There is a problem if the outputs are the same as the inputs # or the first output. # TODO this is a temp fix. - flat_out.append(mx.sym.identity(s)) - g = mx.sym.Group(flat_out) - - # Find free variables in the python that are symbols. - freevars = dict(zip(func.func_code.co_freevars, - (c.cell_contents for c in func.func_closure))) - sym_freevars = [] - for name in freevars: - val = freevars[name] - if isinstance(val, mx.sym.Symbol): - # We need to save the original symbol first. - sym_freevars.append(val) - gin_names.append(name) + flat_out.append(symbol.identity(s)) + g = symbol.Group(flat_out) + input_syms = _get_graph_inputs(g, name, "ro_var") if (isinstance(input, list)): num_inputs = len(input) @@ -61,17 +71,26 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): # Here we need to find out how the input symbols are ordered as well as # where the loop states are located in the list of inputs. - ins = init_states + sym_freevars - ins = {sym.name:sym for sym in ins} + + # This dict contains the symbols of the subgraph. + input_syms = {sym.name:sym for sym in input_syms} + gin_names = input_syms.keys() + # This array contains the symbols for the inputs of foreach. ordered_ins = [] + states_map = {sym.name:sym for sym in init_states} + state_names = states_map.keys() in_state_locs = [-1] * len(init_states) for in_name in g.list_inputs(): - assert in_name in gin_names, "The input graph contains variables we can't find" - if in_name in ins: - ordered_ins.append(ins[in_name]) - for i in range(len(init_states)): - if (init_states[i].name == in_name): - in_state_locs[i] = len(ordered_ins) - 1 + num_inputs + assert in_name in gin_names, "The input variable %s can't be found in graph inputs: %s" \ + % (in_name, str(gin_names)) + if (in_name in state_names): + ordered_ins.append(states_map[in_name]) + elif (in_name != "in"): + ordered_ins.append(input_syms[in_name]) + + for i in range(len(init_states)): + if (init_states[i].name == in_name): + in_state_locs[i] = len(ordered_ins) - 1 + num_inputs - return mx.sym._internal._foreach(g, input, *ordered_ins, num_outputs=len(flat_out), + return symbol._internal._foreach(g, input, *ordered_ins, num_outputs=len(flat_out), in_state_locs=in_state_locs) diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index 6ba8a3d95469..a2ebefc8255d 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -345,6 +345,33 @@ int MXSymbolGetAtomicSymbolName(AtomicSymbolCreator creator, API_END(); } +int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **out_arr, int *out_size) { + API_BEGIN(); + nnvm::Symbol *s = static_cast(sym); + nnvm::Graph g; + g.outputs = s->outputs; + std::vector input_syms; + const nnvm::IndexedGraph& idx = g.indexed_graph(); + size_t max_out_size = *out_size; + // Go through all nodes and return the ones representing variables. + for (size_t i = 0; i < idx.num_nodes(); i++) { + const nnvm::Node &n = *idx[i].source; + for (const nnvm::NodeEntry &e : n.inputs) { + auto p = e.node; + if (p->is_variable()) { + nnvm::Symbol *s = new nnvm::Symbol(); + s->outputs.push_back(e); + input_syms.push_back(s); + std::cout << p->attrs.name << std::endl; + } + } + } + CHECK(input_syms.size() <= max_out_size); + *out_size = input_syms.size(); + memcpy(out_arr, input_syms.data(), sizeof(*out_arr) * input_syms.size()); + API_END(); +} + int MXSymbolCreateFromFile(const char *fname, SymbolHandle *out) { nnvm::Symbol *s = new nnvm::Symbol(); API_BEGIN(); From 409bbe2d5c0e585308c5b2706f7462dc5ffd77e1 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 13 Apr 2018 23:38:52 +0000 Subject: [PATCH 19/71] Fix shape, dtype and storage inference. --- src/operator/nn/control_flow.cc | 36 ++++++++++++++++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 135694742bed..9ada9f105624 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -312,6 +312,16 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, imperative::CheckAndInferShape(g.get(), std::move(shape_inputs), true); const auto& shapes = g->GetAttr("shape"); + // Inferring the shape in the subgraph may infer the shape of the inputs. + // We need to copy the inferred input shapes back. + const auto &input_nids = idx.input_nodes(); + CHECK_EQ(input_nids.size(), in_shape->size()); + size_t num_input_arrays = 1; + for (size_t i = num_input_arrays; i < in_shape->size(); i++) { + auto eid = idx.entry_id(input_nids[i], 0); + (*in_shape)[i] = shapes[eid]; + } + // For the first shape. uint32_t eid = idx.entry_id(g->outputs[0]); const auto& g_out_shape = shapes[eid]; @@ -344,7 +354,19 @@ static bool ForeachType(const nnvm::NodeAttrs& attrs, // TODO(zhengda) This can also be called in the execution engine. // We need to make it thread-safe. imperative::CheckAndInferType(g.get(), std::move(dtype_inputs), true); + + size_t num_input_arrays = 1; const auto &dtypes = g->GetAttr("dtype"); + + // Inferring the data type in the subgraph may infer the data type of the inputs. + // We need to copy the inferred input data types back. + const auto &input_nids = idx.input_nodes(); + CHECK_EQ(input_nids.size(), in_type->size()); + for (size_t i = num_input_arrays; i < in_type->size(); i++) { + auto eid = idx.entry_id(input_nids[i], 0); + (*in_type)[i] = dtypes[eid]; + } + for (size_t i = 0; i < g->outputs.size(); i++) (*out_type)[i] = dtypes[idx.entry_id(g->outputs[i])]; return true; @@ -367,8 +389,20 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, StorageTypeVector storage_type_inputs = *in_attrs; imperative::CheckAndInferStorageType(g.get(), std::move(dev_masks), std::move(storage_type_inputs), true); - *dispatch_mode = DispatchMode::kFComputeEx; + + size_t num_input_arrays = 1; const auto& stypes = g->GetAttr("storage_type"); + + // Inferring the storage in the subgraph may infer the storage of the inputs. + // We need to copy the inferred input storage back. + const auto &input_nids = idx.input_nodes(); + CHECK_EQ(input_nids.size(), in_attrs->size()); + for (size_t i = num_input_arrays; i < in_attrs->size(); i++) { + auto eid = idx.entry_id(input_nids[i], 0); + (*in_attrs)[i] = stypes[eid]; + } + + *dispatch_mode = DispatchMode::kFComputeEx; auto &outputs = idx.outputs(); CHECK(outputs.size() == out_attrs->size()); for (size_t i = 0; i < out_attrs->size(); i++) From 9660340c674fe9fdbca4d47d4a43738588d8f1d0 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 14 Apr 2018 00:02:42 +0000 Subject: [PATCH 20/71] reorganize the output of foreach. --- python/mxnet/contrib/control_flow.py | 17 +++++++++++++++-- tests/python/unittest/test_operator.py | 14 +++++++++++--- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/python/mxnet/contrib/control_flow.py b/python/mxnet/contrib/control_flow.py index 6175df124eb1..df3716040aa9 100644 --- a/python/mxnet/contrib/control_flow.py +++ b/python/mxnet/contrib/control_flow.py @@ -92,5 +92,18 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): if (init_states[i].name == in_name): in_state_locs[i] = len(ordered_ins) - 1 + num_inputs - return symbol._internal._foreach(g, input, *ordered_ins, num_outputs=len(flat_out), - in_state_locs=in_state_locs) + num_outputs = len(flat_out) + num_states = len(state_names) + ret = symbol._internal._foreach(g, input, *ordered_ins, num_outputs=num_outputs, + in_state_locs=in_state_locs) + if (num_outputs - num_states > 1): + outs = [] + for i in range(num_outputs - num_states): + outs.append(ret[i]) + else: + outs = ret[0] + states = [] + for i in range(num_states): + states.append(ret[num_outputs - num_states + i]) + + return (outs, states) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 2904c9c80ded..a14a20d0bd68 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5671,12 +5671,11 @@ def test_foreach(): # This tests foreach with accumulation sum. def step(in1, states): out = in1 + states[0] - # TODO This is problematic. We can't count on the user to define two different symbols. - return (out, [out * 1]) + return (out, [out]) out = mx.contrib.cf.foreach(step, v3, [v4]) out1 = out[0] * 2 - out = mx.sym.Group([out1, out[1]]) + out = mx.sym.Group([out1, out[1][0]]) arr1 = mx.nd.random.uniform(shape=(5, 2)) arr2 = mx.nd.random.uniform(shape=(2)) e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}) @@ -5715,6 +5714,14 @@ def step(in1, states): # TODO This is problematic. We can't count on the user to define two different symbols. return (next_h, [next_h, next_c]) + def sym_group(out): + if (isinstance(out[0], mx.sym.Symbol)): + ret = [out[0]] + else: + ret = out[0] + ret.extend(out[1]) + return mx.sym.Group(ret) + data_arr = mx.nd.random.uniform(shape=(5, 2, 4)) h_arr = mx.nd.random.uniform(shape=(2, 4)) c_arr = mx.nd.random.uniform(shape=(2, 4)) @@ -5724,6 +5731,7 @@ def step(in1, states): h2h_barr = mx.nd.random.uniform(shape=(16)) out = mx.contrib.cf.foreach(step, data, [init_h, init_c]) + out = sym_group(out) e = out.bind(ctx=mx.cpu(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, 'i2h_weight': i2h_warr, 'h2h_weight': h2h_warr, 'i2h_bias': i2h_barr, 'h2h_bias': h2h_barr}) e.forward() From a0f8d5270efd7e1ced485187d19438c9fb36c2f6 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 14 Apr 2018 00:05:09 +0000 Subject: [PATCH 21/71] Add a gluon RNN unroll with symbol foreach. --- python/mxnet/gluon/contrib/rnn/rnn_cell.py | 149 ++++++++++++++++++++- tests/python/unittest/test_gluon_rnn.py | 14 +- 2 files changed, 161 insertions(+), 2 deletions(-) diff --git a/python/mxnet/gluon/contrib/rnn/rnn_cell.py b/python/mxnet/gluon/contrib/rnn/rnn_cell.py index 1b9afee14bf2..7cd2d0a4a98f 100644 --- a/python/mxnet/gluon/contrib/rnn/rnn_cell.py +++ b/python/mxnet/gluon/contrib/rnn/rnn_cell.py @@ -17,10 +17,16 @@ # coding: utf-8 """Definition of various recurrent neural network cells.""" -__all__ = ['VariationalDropoutCell', 'LSTMPCell'] +__all__ = ['VariationalDropoutCell', 'LSTMPCell', 'SymHybridRNNCell', 'RNNCell'] +import inspect + +from .... import symbol, ndarray +from ....base import _as_list +from ....contrib.control_flow import foreach from ...rnn import BidirectionalCell, SequentialRNNCell, ModifierCell, HybridRecurrentCell from ...rnn.rnn_cell import _format_sequence, _get_begin_state, _mask_sequence_variable_length +from ...rnn.rnn_cell import RNNCell as GluonRNNCell from ... import tensor_types class VariationalDropoutCell(ModifierCell): @@ -315,3 +321,144 @@ def hybrid_forward(self, F, inputs, states, i2h_weight, return next_r, [next_r, next_c] # pylint: enable= arguments-differ + +class SymHybridRNNCell(HybridRecurrentCell): + def __init__(self, prefix=None, params=None): + super(SymHybridRNNCell, self).__init__(prefix=prefix, params=params) + + def unroll(self, inputs, begin_state=None, layout='NTC', + merge_outputs=None, valid_length=None): + # if this is a list, we can have unroll in the parent class to handle it. + if (isinstance(inputs, list)): + return super(SymHybridRNNCell, self).unroll(self, len(inputs), inputs, begin_state, + layout, merge_outputs, valid_length) + elif (isinstance(inputs, ndarray.NDArray)): + axis = layout.find('T') + length = inputs.shape[axis] + return super(SymHybridRNNCell, self).unroll(self, length, inputs, begin_state, + layout, merge_outputs, valid_length) + + self.reset() + + batch_size = 0 + F = symbol + axis = layout.find('T') + begin_state = _get_begin_state(self, F, begin_state, inputs, batch_size) + + states = begin_state + outputs = [] + all_states = [] + def iter_func(input, states): + return self(input, states) + outputs, last_states = foreach(iter_func, inputs, begin_state) + #if valid_length is not None: + # states = [F.SequenceLast(ele_list, + # sequence_length=valid_length, + # use_sequence_length=True, + # axis=0) + # for ele_list in all_states] + # outputs = F.SequenceMask(outputs, sequence_length=valid_length, use_sequence_length=True, + # axis=axis) + #outputs, _, _, _ = _format_sequence(length, outputs, layout, merge_outputs) + + return outputs, last_states + +class RNNCell(SymHybridRNNCell): + r"""Elman RNN recurrent neural network cell. + + Each call computes the following function: + + .. math:: + + h_t = \tanh(w_{ih} * x_t + b_{ih} + w_{hh} * h_{(t-1)} + b_{hh}) + + where :math:`h_t` is the hidden state at time `t`, and :math:`x_t` is the hidden + state of the previous layer at time `t` or :math:`input_t` for the first layer. + If nonlinearity='relu', then `ReLU` is used instead of `tanh`. + + Parameters + ---------- + hidden_size : int + Number of units in output symbol + activation : str or Symbol, default 'tanh' + Type of activation function. + i2h_weight_initializer : str or Initializer + Initializer for the input weights matrix, used for the linear + transformation of the inputs. + h2h_weight_initializer : str or Initializer + Initializer for the recurrent weights matrix, used for the linear + transformation of the recurrent state. + i2h_bias_initializer : str or Initializer + Initializer for the bias vector. + h2h_bias_initializer : str or Initializer + Initializer for the bias vector. + prefix : str, default 'rnn_' + Prefix for name of `Block`s + (and name of weight if params is `None`). + params : Parameter or None + Container for weight sharing between cells. + Created if `None`. + + + Inputs: + - **data**: input tensor with shape `(batch_size, input_size)`. + - **states**: a list of one initial recurrent state tensor with shape + `(batch_size, num_hidden)`. + + Outputs: + - **out**: output tensor with shape `(batch_size, num_hidden)`. + - **next_states**: a list of one output recurrent state tensor with the + same shape as `states`. + """ + def __init__(self, hidden_size, activation='tanh', + i2h_weight_initializer=None, h2h_weight_initializer=None, + i2h_bias_initializer='zeros', h2h_bias_initializer='zeros', + input_size=0, prefix=None, params=None): + super(RNNCell, self).__init__(prefix=prefix, params=params) + self._hidden_size = hidden_size + self._activation = activation + self._input_size = input_size + self.i2h_weight = self.params.get('i2h_weight', shape=(hidden_size, input_size), + init=i2h_weight_initializer, + allow_deferred_init=True) + self.h2h_weight = self.params.get('h2h_weight', shape=(hidden_size, hidden_size), + init=h2h_weight_initializer, + allow_deferred_init=True) + self.i2h_bias = self.params.get('i2h_bias', shape=(hidden_size,), + init=i2h_bias_initializer, + allow_deferred_init=True) + self.h2h_bias = self.params.get('h2h_bias', shape=(hidden_size,), + init=h2h_bias_initializer, + allow_deferred_init=True) + + def state_info(self, batch_size=0): + return [{'shape': (batch_size, self._hidden_size), '__layout__': 'NC'}] + + def _alias(self): + return 'rnn' + + def __repr__(self): + s = '{name}({mapping}' + if hasattr(self, '_activation'): + s += ', {_activation}' + s += ')' + shape = self.i2h_weight.shape + mapping = '{0} -> {1}'.format(shape[1] if shape[1] else None, shape[0]) + return s.format(name=self.__class__.__name__, + mapping=mapping, + **self.__dict__) + + def hybrid_forward(self, F, inputs, states, i2h_weight, + h2h_weight, i2h_bias, h2h_bias): + prefix = 't%d_'%self._counter + i2h = F.FullyConnected(data=inputs, weight=i2h_weight, bias=i2h_bias, + num_hidden=self._hidden_size, + name=prefix+'i2h') + h2h = F.FullyConnected(data=states[0], weight=h2h_weight, bias=h2h_bias, + num_hidden=self._hidden_size, + name=prefix+'h2h') + output = self._get_activation(F, i2h + h2h, self._activation, + name=prefix+'out') + + print("contrib.RNNCell") + return output, [output] diff --git a/tests/python/unittest/test_gluon_rnn.py b/tests/python/unittest/test_gluon_rnn.py index 24d5a932d7b2..548789eafdbc 100644 --- a/tests/python/unittest/test_gluon_rnn.py +++ b/tests/python/unittest/test_gluon_rnn.py @@ -28,13 +28,25 @@ def test_rnn(): inputs = [mx.sym.Variable('rnn_t%d_data'%i) for i in range(3)] outputs, _ = cell.unroll(3, inputs) outputs = mx.sym.Group(outputs) - assert sorted(cell.collect_params().keys()) == ['rnn_h2h_bias', 'rnn_h2h_weight', 'rnn_i2h_bias', 'rnn_i2h_weight'] + assert sorted(cell.collect_params().keys()) == ['rnn_h2h_bias', 'rnn_h2h_weight', + 'rnn_i2h_bias', 'rnn_i2h_weight'] assert outputs.list_outputs() == ['rnn_t0_out_output', 'rnn_t1_out_output', 'rnn_t2_out_output'] args, outs, auxs = outputs.infer_shape(rnn_t0_data=(10,50), rnn_t1_data=(10,50), rnn_t2_data=(10,50)) assert outs == [(10, 100), (10, 100), (10, 100)] +def test_contrib_rnn(): + contrib_cell = gluon.contrib.rnn.RNNCell(100, prefix='rnn_') + inputs = mx.sym.Variable('rnn_data') + contrib_outputs, _ = contrib_cell.unroll(inputs) + assert sorted(contrib_cell.collect_params().keys()) == ['rnn_h2h_bias', 'rnn_h2h_weight', + 'rnn_i2h_bias', 'rnn_i2h_weight'] + + args, outs, auxs = contrib_outputs.infer_shape(rnn_data=(3, 10,50)) + assert outs == [(3, 10, 100)] + + def test_lstm(): cell = gluon.rnn.LSTMCell(100, prefix='rnn_') inputs = [mx.sym.Variable('rnn_t%d_data'%i) for i in range(3)] From f462cc700537791e1afee95b9e0657e249169ccd Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 16 Apr 2018 21:13:47 +0000 Subject: [PATCH 22/71] print unnecessary print. --- python/mxnet/gluon/contrib/rnn/rnn_cell.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/mxnet/gluon/contrib/rnn/rnn_cell.py b/python/mxnet/gluon/contrib/rnn/rnn_cell.py index 7cd2d0a4a98f..27871a56e315 100644 --- a/python/mxnet/gluon/contrib/rnn/rnn_cell.py +++ b/python/mxnet/gluon/contrib/rnn/rnn_cell.py @@ -460,5 +460,4 @@ def hybrid_forward(self, F, inputs, states, i2h_weight, output = self._get_activation(F, i2h + h2h, self._activation, name=prefix+'out') - print("contrib.RNNCell") return output, [output] From 93280346a078543dd63473d8208141369d7646a6 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 16 Apr 2018 23:26:26 +0000 Subject: [PATCH 23/71] have imperative and symbolic foreach. --- python/mxnet/contrib/control_flow.py | 109 --------------------- python/mxnet/gluon/contrib/rnn/rnn_cell.py | 18 ++-- python/mxnet/ndarray/contrib.py | 15 +++ python/mxnet/symbol/contrib.py | 93 ++++++++++++++++++ tests/python/unittest/test_gluon_rnn.py | 18 ++++ tests/python/unittest/test_operator.py | 4 +- 6 files changed, 136 insertions(+), 121 deletions(-) delete mode 100644 python/mxnet/contrib/control_flow.py diff --git a/python/mxnet/contrib/control_flow.py b/python/mxnet/contrib/control_flow.py deleted file mode 100644 index df3716040aa9..000000000000 --- a/python/mxnet/contrib/control_flow.py +++ /dev/null @@ -1,109 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -import ctypes - -from .. import symbol -from ..base import _LIB, c_str, c_array, check_call -from ..base import SymbolHandle, NDArrayHandle -from ..attribute import AttrScope - -def _get_graph_inputs(subg, name, prefix): - num_handles = ctypes.c_int(1000) - handles = c_array(SymbolHandle, [SymbolHandle(0) for i in range(1000)]) - check_call(_LIB.MXSymbolGetInputSymbols(subg.handle, handles, - ctypes.byref(num_handles))) - - syms = [] - for i in range(num_handles.value): - s = symbol.Symbol(handles[i]) - syms.append(s) - return syms - -def foreach(func, input, init_states, back_prop=False, name="foreach"): - assert isinstance(init_states, list), "init_states should be a list" - states = [] - with AttrScope(subgraph_name=name): - in_ele = symbol.var("in") - for s in init_states: - states.append(symbol.var(s.name)) - - sym_out = func(in_ele, states) - # The function should return a tuple. The first element goes to - # the output of the function. The second element is a list. - assert isinstance(sym_out, tuple), "func should return a tuple (out, states)" - assert isinstance(sym_out[1], list), \ - "the second element in the returned tuple should be a list" - assert len(sym_out[1]) == len(init_states), \ - "the number of output states (%d) should be the same as input states (%d)" \ - % (len(sym_out[1]), len(init_states)) - - if (isinstance(sym_out[0], list)): - flat_out = sym_out[0] - else: - flat_out = [sym_out[0]] - for s in sym_out[1]: - # There is a problem if the outputs are the same as the inputs - # or the first output. - # TODO this is a temp fix. - flat_out.append(symbol.identity(s)) - g = symbol.Group(flat_out) - input_syms = _get_graph_inputs(g, name, "ro_var") - - if (isinstance(input, list)): - num_inputs = len(input) - else: - num_inputs = 1 - - # Here we need to find out how the input symbols are ordered as well as - # where the loop states are located in the list of inputs. - - # This dict contains the symbols of the subgraph. - input_syms = {sym.name:sym for sym in input_syms} - gin_names = input_syms.keys() - # This array contains the symbols for the inputs of foreach. - ordered_ins = [] - states_map = {sym.name:sym for sym in init_states} - state_names = states_map.keys() - in_state_locs = [-1] * len(init_states) - for in_name in g.list_inputs(): - assert in_name in gin_names, "The input variable %s can't be found in graph inputs: %s" \ - % (in_name, str(gin_names)) - if (in_name in state_names): - ordered_ins.append(states_map[in_name]) - elif (in_name != "in"): - ordered_ins.append(input_syms[in_name]) - - for i in range(len(init_states)): - if (init_states[i].name == in_name): - in_state_locs[i] = len(ordered_ins) - 1 + num_inputs - - num_outputs = len(flat_out) - num_states = len(state_names) - ret = symbol._internal._foreach(g, input, *ordered_ins, num_outputs=num_outputs, - in_state_locs=in_state_locs) - if (num_outputs - num_states > 1): - outs = [] - for i in range(num_outputs - num_states): - outs.append(ret[i]) - else: - outs = ret[0] - states = [] - for i in range(num_states): - states.append(ret[num_outputs - num_states + i]) - - return (outs, states) diff --git a/python/mxnet/gluon/contrib/rnn/rnn_cell.py b/python/mxnet/gluon/contrib/rnn/rnn_cell.py index 27871a56e315..98b89d22d5c3 100644 --- a/python/mxnet/gluon/contrib/rnn/rnn_cell.py +++ b/python/mxnet/gluon/contrib/rnn/rnn_cell.py @@ -23,7 +23,6 @@ from .... import symbol, ndarray from ....base import _as_list -from ....contrib.control_flow import foreach from ...rnn import BidirectionalCell, SequentialRNNCell, ModifierCell, HybridRecurrentCell from ...rnn.rnn_cell import _format_sequence, _get_begin_state, _mask_sequence_variable_length from ...rnn.rnn_cell import RNNCell as GluonRNNCell @@ -332,17 +331,16 @@ def unroll(self, inputs, begin_state=None, layout='NTC', if (isinstance(inputs, list)): return super(SymHybridRNNCell, self).unroll(self, len(inputs), inputs, begin_state, layout, merge_outputs, valid_length) - elif (isinstance(inputs, ndarray.NDArray)): - axis = layout.find('T') - length = inputs.shape[axis] - return super(SymHybridRNNCell, self).unroll(self, length, inputs, begin_state, - layout, merge_outputs, valid_length) self.reset() - - batch_size = 0 - F = symbol + batch_axis = layout.find('N') axis = layout.find('T') + batch_size = 0 + if isinstance(inputs, symbol.Symbol): + F = symbol + else: + batch_size = inputs.shape[batch_axis] + F = ndarray begin_state = _get_begin_state(self, F, begin_state, inputs, batch_size) states = begin_state @@ -350,7 +348,7 @@ def unroll(self, inputs, begin_state=None, layout='NTC', all_states = [] def iter_func(input, states): return self(input, states) - outputs, last_states = foreach(iter_func, inputs, begin_state) + outputs, last_states = F.contrib.foreach(iter_func, inputs, begin_state) #if valid_length is not None: # states = [F.SequenceLast(ele_list, # sequence_length=valid_length, diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py index ba402e6f3f8d..2d0567eaa9b2 100644 --- a/python/mxnet/ndarray/contrib.py +++ b/python/mxnet/ndarray/contrib.py @@ -96,3 +96,18 @@ def rand_zipfian(true_classes, num_sampled, range_max, ctx=None): expected_count_sampled = expected_prob_sampled * num_sampled return sampled_classes, expected_count_true, expected_count_sampled # pylint: enable=line-too-long + +def foreach(func, input, init_states, back_prop=False, name="foreach"): + assert isinstance(init_states, list), "init_states should be a list" + states = init_states + outputs = [] + for i in range(input.shape[0]): + ele = input[i] + outs, states = func(ele, states) + outs = _as_list(outs) + if (i == 0): + outputs = outs + else: + for j in range(outs): + outputs[j].append(outs[j]) + return (outputs, states) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index 83e90e687327..856591ca0c09 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -26,6 +26,13 @@ except ImportError: pass +import ctypes + +from . import symbol +from ..base import _LIB, c_str, c_array, check_call +from ..base import SymbolHandle, NDArrayHandle +from ..attribute import AttrScope + __all__ = ["rand_zipfian"] def rand_zipfian(true_classes, num_sampled, range_max): @@ -91,3 +98,89 @@ def rand_zipfian(true_classes, num_sampled, range_max): expected_prob_sampled = ((sampled_cls_fp64 + 2.0) / (sampled_cls_fp64 + 1.0)).log() / log_range expected_count_sampled = expected_prob_sampled * num_sampled return sampled_classes, expected_count_true, expected_count_sampled + +def _get_graph_inputs(subg, name, prefix): + num_handles = ctypes.c_int(1000) + handles = c_array(SymbolHandle, [SymbolHandle(0) for i in range(1000)]) + check_call(_LIB.MXSymbolGetInputSymbols(subg.handle, handles, + ctypes.byref(num_handles))) + + syms = [] + for i in range(num_handles.value): + s = Symbol(handles[i]) + syms.append(s) + return syms + +def foreach(func, input, init_states, back_prop=False, name="foreach"): + assert isinstance(init_states, list), "init_states should be a list" + states = [] + with AttrScope(subgraph_name=name): + in_ele = symbol.var("in") + for s in init_states: + states.append(symbol.var(s.name)) + + sym_out = func(in_ele, states) + # The function should return a tuple. The first element goes to + # the output of the function. The second element is a list. + assert isinstance(sym_out, tuple), "func should return a tuple (out, states)" + assert isinstance(sym_out[1], list), \ + "the second element in the returned tuple should be a list" + assert len(sym_out[1]) == len(init_states), \ + "the number of output states (%d) should be the same as input states (%d)" \ + % (len(sym_out[1]), len(init_states)) + + if (isinstance(sym_out[0], list)): + flat_out = sym_out[0] + else: + flat_out = [sym_out[0]] + for s in sym_out[1]: + # There is a problem if the outputs are the same as the inputs + # or the first output. + # TODO this is a temp fix. + flat_out.append(symbol.op.identity(s)) + g = symbol.Group(flat_out) + input_syms = _get_graph_inputs(g, name, "ro_var") + + if (isinstance(input, list)): + num_inputs = len(input) + else: + num_inputs = 1 + + # Here we need to find out how the input symbols are ordered as well as + # where the loop states are located in the list of inputs. + + # This dict contains the symbols of the subgraph. + input_syms = {sym.name:sym for sym in input_syms} + gin_names = input_syms.keys() + # This array contains the symbols for the inputs of foreach. + ordered_ins = [] + states_map = {sym.name:sym for sym in init_states} + state_names = states_map.keys() + in_state_locs = [-1] * len(init_states) + for in_name in g.list_inputs(): + assert in_name in gin_names, "The input variable %s can't be found in graph inputs: %s" \ + % (in_name, str(gin_names)) + if (in_name in state_names): + ordered_ins.append(states_map[in_name]) + elif (in_name != "in"): + ordered_ins.append(input_syms[in_name]) + + for i in range(len(init_states)): + if (init_states[i].name == in_name): + in_state_locs[i] = len(ordered_ins) - 1 + num_inputs + + num_outputs = len(flat_out) + num_states = len(state_names) + ret = symbol._internal._foreach(g, input, *ordered_ins, num_outputs=num_outputs, + in_state_locs=in_state_locs) + if (num_outputs - num_states > 1): + outs = [] + for i in range(num_outputs - num_states): + outs.append(ret[i]) + else: + outs = ret[0] + states = [] + for i in range(num_states): + states.append(ret[num_outputs - num_states + i]) + + return (outs, states) diff --git a/tests/python/unittest/test_gluon_rnn.py b/tests/python/unittest/test_gluon_rnn.py index 548789eafdbc..cd885695412f 100644 --- a/tests/python/unittest/test_gluon_rnn.py +++ b/tests/python/unittest/test_gluon_rnn.py @@ -36,6 +36,14 @@ def test_rnn(): assert outs == [(10, 100), (10, 100), (10, 100)] +class RNNLayer(gluon.HybridBlock): + def __init__(self, prefix=None, params=None): + super(RNNLayer, self).__init__(prefix=prefix, params=params) + self.cell = gluon.contrib.rnn.RNNCell(100, prefix='rnn_') + + def hybrid_forward(self, F, inputs, states=None): + return self.cell.unroll(inputs, states) + def test_contrib_rnn(): contrib_cell = gluon.contrib.rnn.RNNCell(100, prefix='rnn_') inputs = mx.sym.Variable('rnn_data') @@ -46,6 +54,16 @@ def test_contrib_rnn(): args, outs, auxs = contrib_outputs.infer_shape(rnn_data=(3, 10,50)) assert outs == [(3, 10, 100)] + rnn_data = mx.nd.normal(loc=0, scale=1, shape=(3, 10, 50)) + layer = RNNLayer() + layer.initialize(ctx=mx.cpu(0)) + res1 = layer(rnn_data) + + layer = RNNLayer() + layer.initialize(ctx=mx.cpu(0)) + layer.hybridize() + res2 = layer(rnn_data) + def test_lstm(): cell = gluon.rnn.LSTMCell(100, prefix='rnn_') diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index a14a20d0bd68..38507771407a 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5673,7 +5673,7 @@ def step(in1, states): out = in1 + states[0] return (out, [out]) - out = mx.contrib.cf.foreach(step, v3, [v4]) + out = mx.sym.contrib.foreach(step, v3, [v4]) out1 = out[0] * 2 out = mx.sym.Group([out1, out[1][0]]) arr1 = mx.nd.random.uniform(shape=(5, 2)) @@ -5730,7 +5730,7 @@ def sym_group(out): i2h_barr = mx.nd.random.uniform(shape=(16)) h2h_barr = mx.nd.random.uniform(shape=(16)) - out = mx.contrib.cf.foreach(step, data, [init_h, init_c]) + out = mx.sym.contrib.foreach(step, data, [init_h, init_c]) out = sym_group(out) e = out.bind(ctx=mx.cpu(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, 'i2h_weight': i2h_warr, 'h2h_weight': h2h_warr, 'i2h_bias': i2h_barr, 'h2h_bias': h2h_barr}) From 62d767f4c74d93c12c26864ad0067089d64b2066 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 18 Apr 2018 01:57:42 +0000 Subject: [PATCH 24/71] Fix an error after moving foreach. --- python/mxnet/contrib/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/mxnet/contrib/__init__.py b/python/mxnet/contrib/__init__.py index 7489d97d90fe..fbfd3469678b 100644 --- a/python/mxnet/contrib/__init__.py +++ b/python/mxnet/contrib/__init__.py @@ -32,4 +32,3 @@ from . import io from . import quantization from . import quantization as quant -from . import control_flow as cf From f8c4383fe3ff62a9109e119bebe1a87712711a0e Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 18 Apr 2018 01:58:28 +0000 Subject: [PATCH 25/71] Fix imperative foreach --- python/mxnet/ndarray/contrib.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py index 2d0567eaa9b2..7bf96d8f68f0 100644 --- a/python/mxnet/ndarray/contrib.py +++ b/python/mxnet/ndarray/contrib.py @@ -21,6 +21,8 @@ import math from ..context import current_context from ..random import uniform +from ..base import _as_list +from .op import stack try: from .gen_contrib import * except ImportError: @@ -106,8 +108,12 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): outs, states = func(ele, states) outs = _as_list(outs) if (i == 0): - outputs = outs + # outputs is a list of lists + for j in range(len(outs)): + outputs.append([outs[j]]) else: - for j in range(outs): + for j in range(len(outs)): outputs[j].append(outs[j]) + for i in range(len(outputs)): + outputs[i] = stack(*outputs[i]) return (outputs, states) From 15d3619c235b5746bbf770a7d94b4b3baa3a0455 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 24 Apr 2018 02:10:42 +0000 Subject: [PATCH 26/71] Fix a minor problem. --- python/mxnet/gluon/contrib/rnn/rnn_cell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mxnet/gluon/contrib/rnn/rnn_cell.py b/python/mxnet/gluon/contrib/rnn/rnn_cell.py index 98b89d22d5c3..dcb396a57613 100644 --- a/python/mxnet/gluon/contrib/rnn/rnn_cell.py +++ b/python/mxnet/gluon/contrib/rnn/rnn_cell.py @@ -325,7 +325,7 @@ class SymHybridRNNCell(HybridRecurrentCell): def __init__(self, prefix=None, params=None): super(SymHybridRNNCell, self).__init__(prefix=prefix, params=params) - def unroll(self, inputs, begin_state=None, layout='NTC', + def unroll(self, inputs, begin_state=None, layout='TNC', merge_outputs=None, valid_length=None): # if this is a list, we can have unroll in the parent class to handle it. if (isinstance(inputs, list)): From 191fec203f0ce575de412fef919c3a84fa5e0ecb Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 30 Apr 2018 18:56:46 +0000 Subject: [PATCH 27/71] Use CachedOp to execute subgraph. --- src/operator/nn/control_flow.cc | 174 ++++---------------------------- 1 file changed, 18 insertions(+), 156 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 9ada9f105624..fd83d3b07363 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -30,132 +30,21 @@ namespace mxnet { namespace op { -void RunGraph(const nnvm::IndexedGraph& idx, - const std::vector arrays, - size_t node_start, size_t node_end, - std::vector&& array_reqs, - std::vector&& ref_count, - std::vector *p_states, - const DispatchModeVector &dispatch_modes) { - using namespace nnvm; - using namespace imperative; - static auto& createop = nnvm::Op::GetAttr("FCreateOpState"); - static auto& is_layer_backward = Op::GetAttr("TIsLayerOpBackward"); - - std::vector& states = *p_states; - std::vector ndinputs, ndoutputs; - ShapeVector arg_shapes; - DTypeVector arg_dtypes; - std::vector req; - - for (size_t i = node_start; i < node_end; ++i) { - const nnvm::IndexedGraph::Node& node = idx[i]; - if (node.source->op() == nullptr) continue; - auto num_outputs = node.source->num_outputs(); - ndinputs.clear(); - ndinputs.reserve(node.inputs.size()); - for (const auto& j : node.inputs) { - ndinputs.emplace_back(arrays[idx.entry_id(j)]); - CHECK(!ndinputs.back()->is_none()) << idx[j.node_id].source->attrs.name - << " " << j.index; - } - ndoutputs.clear(); - ndoutputs.reserve(num_outputs); - req.clear(); - req.reserve(num_outputs); - for (size_t j = 0; j < num_outputs; ++j) { - size_t eid = idx.entry_id(i, j); - ndoutputs.emplace_back(arrays[eid]); - req.push_back(array_reqs[eid]); - CHECK(!ndoutputs.back()->is_none()); - } - const Context& ctx = ndoutputs[0]->ctx(); - const DispatchMode dispatch_mode = dispatch_modes[i]; - if (createop.count(node.source->op())) { - arg_shapes.clear(); - arg_dtypes.clear(); - arg_shapes.reserve(ndinputs.size()); - arg_dtypes.reserve(ndinputs.size()); - for (size_t i = 0; i < ndinputs.size(); ++i) { - arg_shapes.emplace_back(ndinputs[i]->shape()); - arg_dtypes.emplace_back(ndinputs[i]->dtype()); - } - states[i] = createop[node.source->op()]( - node.source->attrs, ctx, arg_shapes, arg_dtypes); - Imperative::InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, req, - dispatch_mode, states[i]); - } else if (is_layer_backward.get(node.source->op(), false)) { - nnvm::Node* fwd_node = node.source->control_deps[0].get(); - auto fwd_node_id = idx.node_id(fwd_node); - Imperative::InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, - req, dispatch_mode, states[fwd_node_id]); - } else { - Imperative::InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, - req, dispatch_mode); - } - } -} - -static void ExecSubgraph(nnvm::Graph &g, const OpContext& ctx, - const std::vector& cinputs, +static void ExecSubgraph(nnvm::Symbol &sym, const OpContext& ctx, + std::vector cinputs, const std::vector& req, - const std::vector& coutputs) { + std::vector coutputs) { using namespace nnvm; using namespace imperative; - const auto& idx = g.indexed_graph(); - size_t num_inputs = idx.input_nodes().size(); - - CHECK_EQ(num_inputs, cinputs.size()) - << "The subgraph requires " << num_inputs << " but got " << cinputs.size(); - Context default_ctx = cinputs[0].ctx(); - for (size_t i = 0; i < cinputs.size(); ++i) { - CHECK_EQ(cinputs[i].ctx(), default_ctx) - << "The subgraph requires all inputs to live on the same context. But " - << idx[idx.input_nodes()[0]].source->attrs.name << " is on " << default_ctx - << " while " << idx[idx.input_nodes()[i]].source->attrs.name << " is on " - << cinputs[i].ctx(); - } - - // TODO(zhengda) we might want to buffer them. - std::vector buff; - std::vector states; - std::vector inputs = cinputs; - std::vector outputs = coutputs; - - // Allocate entries - states.resize(idx.num_nodes()); - buff.resize(idx.num_node_entries()); - states.reserve(idx.num_nodes()); - std::vector arrays; - arrays.reserve(buff.size()); - for (size_t i = 0; i < buff.size(); ++i) arrays.push_back(&buff[i]); - for (size_t i = 0; i < num_inputs; ++i) { - arrays[idx.entry_id(idx.input_nodes()[i], 0)] = &inputs[i]; - } - for (size_t i = 0; i < idx.outputs().size(); ++i) { - auto eid = idx.entry_id(idx.outputs()[i]); - if (!arrays[eid]->is_none()) outputs[i] = arrays[eid]->Detach(); - arrays[eid] = &outputs[i]; - } - - // Allocate memory for the NDArrays - std::vector ref_count = g.GetAttr >( - ctx.is_train ? "full_ref_count" : "forward_ref_count"); - - std::vector array_reqs(arrays.size(), kWriteTo); - for (size_t i = 0; i < idx.num_node_entries(); ++i) { - if (ref_count[i] == 0) array_reqs[i] = kNullOp; - } - - const auto& mem_plan = g.GetAttr( - ctx.is_train ? "full_mem_plan" : "forward_mem_plan"); - AllocateMemory(g, idx, default_ctx, 0, idx.num_node_entries(), - mem_plan, arrays, &array_reqs); - - const auto& dispatch_modes = g.GetAttr("dispatch_mode"); - RunGraph(idx, arrays, 0, idx.num_nodes(), std::move(array_reqs), - std::move(ref_count), &states, dispatch_modes); + std::vector inputs(cinputs.size()); + std::vector outputs(coutputs.size()); + for (size_t i = 0; i < inputs.size(); i++) + inputs[i] = &cinputs[i]; + for (size_t i = 0; i < outputs.size(); i++) + outputs[i] = &coutputs[i]; + Imperative::CachedOp op(sym, std::vector >()); + op.Forward(nullptr, inputs, outputs); } struct ForeachParam : public dmlc::Parameter { @@ -205,33 +94,6 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, const ForeachParam& params = nnvm::get(attrs.parsed); CHECK_EQ(outputs.size(), (size_t) params.num_outputs); CHECK_EQ(attrs.subgraphs.size(), 1U); - nnvm::Graph &g = *attrs.subgraphs[0]; - const auto& idx = g.indexed_graph(); - - // If this is inference, we only need the forward memory plan. - bool has_mem_plan = !ctx.is_train && g.attrs.count("forward_mem_plan"); - // If this is training, we need the full memory plan. - has_mem_plan = has_mem_plan || (ctx.is_train && g.attrs.count("full_mem_plan")); - // If we don't have a memory plan yet, we need to create a memory plan. - if (!has_mem_plan) { - nnvm::StorageVector storage(idx.num_node_entries(), exec::kBadStorageID); - for (const auto i : idx.input_nodes()) - storage[idx.entry_id(i, 0)] = exec::kExternalStorageID; - const auto& stypes = g.GetAttr("storage_type"); - CHECK_EQ(stypes.size(), storage.size()); - for (size_t i = 0; i < stypes.size(); i++) { - if (stypes[i] != kDefaultStorage) - storage[i] = exec::kDynamicStorageID; - } - - auto mem_plan = imperative::PlanMemory( - &g, std::move(storage), g.GetAttr >( - ctx.is_train ? "full_ref_count" : "forward_ref_count")); - // TODO(zhengda) we need to be careful of changing graph attributes. - // It's not thread-safe. - g.attrs[ctx.is_train ? "full_mem_plan" : "forward_mem_plan"] - = std::make_shared(std::move(mem_plan)); - } size_t len = inputs[0].shape()[0]; CHECK_EQ(inputs[0].shape()[0], outputs[0].shape()[0]); @@ -284,7 +146,7 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, } } - ExecSubgraph(g, ctx, subg_inputs, req, *subg_out_curr); + ExecSubgraph(*attrs.subgraphs[0], ctx, subg_inputs, req, *subg_out_curr); // We need to wait for the iteration to complete before executing // the next one or return from the loop. In this way, we can reuse // the memory in the subgraph. @@ -302,8 +164,8 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, // foreach iterates over the first input NDArray over the first dimension. shape_inputs[0] = TShape(in_shape->at(0).begin() + 1, in_shape->at(0).end()); CHECK_EQ(attrs.subgraphs.size(), 1U); - auto g = attrs.subgraphs[0]; - CHECK(g); + auto g = std::make_shared(); + g->outputs = attrs.subgraphs[0]->outputs; const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_shape->size()); CHECK_EQ(idx.outputs().size(), out_shape->size()); @@ -346,8 +208,8 @@ static bool ForeachType(const nnvm::NodeAttrs& attrs, CHECK_EQ(out_type->size(), (size_t) params.num_outputs); nnvm::DTypeVector dtype_inputs = *in_type; CHECK_EQ(attrs.subgraphs.size(), 1U); - auto g = attrs.subgraphs[0]; - CHECK(g); + auto g = std::make_shared(); + g->outputs = attrs.subgraphs[0]->outputs; const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_type->size()); CHECK_EQ(idx.outputs().size(), out_type->size()); @@ -380,8 +242,8 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, const ForeachParam& params = nnvm::get(attrs.parsed); CHECK_EQ(out_attrs->size(), (size_t) params.num_outputs); CHECK_EQ(attrs.subgraphs.size(), 1U); - auto g = attrs.subgraphs[0]; - CHECK(g); + auto g = std::make_shared(); + g->outputs = attrs.subgraphs[0]->outputs; const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_attrs->size()); CHECK_EQ(idx.outputs().size(), out_attrs->size()); From 65f98e9a66db38308f99b105c10570f9b814d5be Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 1 May 2018 18:21:40 +0000 Subject: [PATCH 28/71] update TODO. --- src/operator/nn/control_flow.cc | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index fd83d3b07363..4f76ccae139d 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -166,11 +166,10 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, CHECK_EQ(attrs.subgraphs.size(), 1U); auto g = std::make_shared(); g->outputs = attrs.subgraphs[0]->outputs; + // TODO(zhengda) We should avoid creating an index graph so many times. const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_shape->size()); CHECK_EQ(idx.outputs().size(), out_shape->size()); - // TODO(zhengda) This can also be called in the execution engine. - // We need to make it thread-safe. imperative::CheckAndInferShape(g.get(), std::move(shape_inputs), true); const auto& shapes = g->GetAttr("shape"); @@ -210,11 +209,10 @@ static bool ForeachType(const nnvm::NodeAttrs& attrs, CHECK_EQ(attrs.subgraphs.size(), 1U); auto g = std::make_shared(); g->outputs = attrs.subgraphs[0]->outputs; + // TODO(zhengda) We should avoid creating an index graph so many times. const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_type->size()); CHECK_EQ(idx.outputs().size(), out_type->size()); - // TODO(zhengda) This can also be called in the execution engine. - // We need to make it thread-safe. imperative::CheckAndInferType(g.get(), std::move(dtype_inputs), true); size_t num_input_arrays = 1; @@ -244,6 +242,7 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, CHECK_EQ(attrs.subgraphs.size(), 1U); auto g = std::make_shared(); g->outputs = attrs.subgraphs[0]->outputs; + // TODO(zhengda) We should avoid creating an index graph so many times. const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_attrs->size()); CHECK_EQ(idx.outputs().size(), out_attrs->size()); From 5e9ec40718568d452cf0e1fa077a5b05140fccf9 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 1 May 2018 23:33:59 +0000 Subject: [PATCH 29/71] make foreach op use FStatefulComputeEx. --- src/executor/attach_op_execs_pass.cc | 30 +++++++++++++++++++--------- src/operator/nn/control_flow.cc | 29 ++++++++++++++++++++++----- 2 files changed, 45 insertions(+), 14 deletions(-) diff --git a/src/executor/attach_op_execs_pass.cc b/src/executor/attach_op_execs_pass.cc index 1944927d66a9..06b01b4f67e3 100644 --- a/src/executor/attach_op_execs_pass.cc +++ b/src/executor/attach_op_execs_pass.cc @@ -126,6 +126,10 @@ class StatefulComputeExecutor : public StorageFallbackOpExecutor { PostFCompute(is_gpu); } + bool HasSubgraph() const override { + return !attrs_.subgraphs.empty(); + } + ExecType exec_type() const override { return exec_type_; } @@ -134,15 +138,16 @@ class StatefulComputeExecutor : public StorageFallbackOpExecutor { return state_.get_var(); } - explicit StatefulComputeExecutor(const OpStatePtr& state, + explicit StatefulComputeExecutor(const NodeAttrs& attrs, const OpStatePtr& state, const FStatefulCompute& fcompute, ExecType exec_type, const std::vector &mutate_idx) - : StorageFallbackOpExecutor(mutate_idx), + : StorageFallbackOpExecutor(mutate_idx), attrs_(attrs), state_(state), fcompute_(fcompute), exec_type_(exec_type) {} private: friend Graph AttachOpExecs(Graph g); + NodeAttrs attrs_; OpStatePtr state_; FStatefulCompute fcompute_; ExecType exec_type_; @@ -160,6 +165,10 @@ class StatefulComputeExExecutor : public OpExecutor { fcompute_(state_, op_ctx, in_array, req, out_array); } + bool HasSubgraph() const override { + return !attrs_.subgraphs.empty(); + } + void Setup() override {} ExecType exec_type() const override { @@ -170,13 +179,14 @@ class StatefulComputeExExecutor : public OpExecutor { return state_.get_var(); } - explicit StatefulComputeExExecutor(const OpStatePtr& state, + explicit StatefulComputeExExecutor(const NodeAttrs& attrs, const OpStatePtr& state, const FStatefulComputeEx& fcompute, ExecType exec_type) - : state_(state), fcompute_(fcompute), exec_type_(exec_type) {} + : attrs_(attrs), state_(state), fcompute_(fcompute), exec_type_(exec_type) {} private: friend Graph AttachOpExecs(Graph g); + NodeAttrs attrs_; OpStatePtr state_; FStatefulComputeEx fcompute_; ExecType exec_type_; @@ -297,15 +307,17 @@ Graph AttachOpExecs(Graph g) { op, "FStatefulComputeEx", vctx[i]); // FStatefulComputeEx is dispatched only when dispatch_mode is DispatchMode::kFComputeEx if (fcompute_ex != nullptr && dispatch_modes[i] == DispatchMode::kFComputeEx) { - ret[i] = std::make_shared(state, fcompute_ex, exec_type); + ret[i] = std::make_shared(inode.source->attrs, state, + fcompute_ex, exec_type); } else { FStatefulCompute fcompute = common::GetFCompute( op, "FStatefulCompute", vctx[i]); CHECK(fcompute != nullptr) << "One of FStatefulCompute and FStatefulComputeEx must be registered " << "for stateful operator " << op->name; - ret[i] = std::make_shared(state, fcompute, - exec_type, mutate_index); + ret[i] = std::make_shared(inode.source->attrs, state, + fcompute, exec_type, + mutate_index); } } else if (is_layer_backward.get(op, false)) { CHECK_GE(inode.control_deps.size(), 1); @@ -316,7 +328,7 @@ Graph AttachOpExecs(Graph g) { op, "FStatefulComputeEx", vctx[i]); // FStatefulComputeEx is dispatched only when dispatch_mode is DispatchMode::kFComputeEx if (fcompute_ex != nullptr && dispatch_modes[i] == DispatchMode::kFComputeEx) { - ret[i] = std::make_shared( + ret[i] = std::make_shared(inode.source->attrs, dynamic_cast(ret[fwd_id].get())->state_, fcompute_ex, exec_type); } else { @@ -325,7 +337,7 @@ Graph AttachOpExecs(Graph g) { CHECK(fcompute != nullptr) << "One of FStatefulCompute and FStatefulComputeEx must be registered " << "for stateful operator " << op->name; - ret[i] = std::make_shared( + ret[i] = std::make_shared(inode.source->attrs, dynamic_cast(ret[fwd_id].get())->state_, fcompute, exec_type, mutate_index); } diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 4f76ccae139d..965dfa1b6560 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -86,14 +86,24 @@ static std::vector ReorderInputs(const std::vector &in, const nnvm::Indexe return ret; } -static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, +struct ForeachState { + Symbol subgraph; + ForeachParam params; + + ForeachState(const Symbol &g, const ForeachParam ¶ms) { + this->subgraph = g; + this->params = params; + } +}; + +static void ForeachComputeExCPU(const OpStatePtr& state_ptr, const OpContext& ctx, const std::vector& inputs, const std::vector& req, const std::vector& outputs) { - const ForeachParam& params = nnvm::get(attrs.parsed); + ForeachState state = state_ptr.get_state(); + const ForeachParam& params = state.params; CHECK_EQ(outputs.size(), (size_t) params.num_outputs); - CHECK_EQ(attrs.subgraphs.size(), 1U); size_t len = inputs[0].shape()[0]; CHECK_EQ(inputs[0].shape()[0], outputs[0].shape()[0]); @@ -146,7 +156,7 @@ static void ForeachComputeExCPU(const nnvm::NodeAttrs& attrs, } } - ExecSubgraph(*attrs.subgraphs[0], ctx, subg_inputs, req, *subg_out_curr); + ExecSubgraph(state.subgraph, ctx, subg_inputs, req, *subg_out_curr); // We need to wait for the iteration to complete before executing // the next one or return from the loop. In this way, we can reuse // the memory in the subgraph. @@ -271,6 +281,14 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, return true; } +OpStatePtr CreateForeachState(const NodeAttrs& attrs, + Context ctx, + const std::vector& ishape, + const std::vector& itype) { + const ForeachParam& params = nnvm::get(attrs.parsed); + return OpStatePtr::Create(*attrs.subgraphs[0], params); +} + NNVM_REGISTER_OP(_foreach) .describe(R"code(foreach)code" ADD_FILELINE) .set_attr_parser(ParamParser) @@ -291,9 +309,10 @@ NNVM_REGISTER_OP(_foreach) [](const NodeAttrs& attrs) { return std::vector{0}; }) +.set_attr("FCreateOpState", CreateForeachState) .set_attr("FInferShape", ForeachShape) .set_attr("FInferType", ForeachType) -.set_attr("FComputeEx", ForeachComputeExCPU) +.set_attr("FStatefulComputeEx", ForeachComputeExCPU) .set_attr("key_var_num_args", "num_args") .add_argument("fn", "Symbol", "Input graph.") .add_argument("input", "NDArray-or-Symbol", "The input array where we iterate over.") From f3ce49c79a7ca78cfd0152379f41da17ba7318ce Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 2 May 2018 22:08:23 +0000 Subject: [PATCH 30/71] Add backward. --- include/mxnet/ndarray.h | 4 + src/operator/nn/control_flow.cc | 222 ++++++++++++++++++++++--- tests/python/unittest/test_operator.py | 41 +++-- 3 files changed, 230 insertions(+), 37 deletions(-) diff --git a/include/mxnet/ndarray.h b/include/mxnet/ndarray.h index e243eb71c477..c6f181c120d2 100644 --- a/include/mxnet/ndarray.h +++ b/include/mxnet/ndarray.h @@ -693,6 +693,10 @@ class NDArray { NDArray MKLDNNDataReshape(const TShape &shape) const; #endif + const nnvm::NodeEntry &GetAutogradEntry() const { + return entry_; + } + /*! * \brief Save list of ndarray into the Stream.x * \param fo The stream of output. diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 965dfa1b6560..e84b35267cb8 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -25,28 +25,12 @@ #include #include #include "../operator_common.h" +#include "../elemwise_op_common.h" #include "../../imperative/imperative_utils.h" namespace mxnet { namespace op { -static void ExecSubgraph(nnvm::Symbol &sym, const OpContext& ctx, - std::vector cinputs, - const std::vector& req, - std::vector coutputs) { - using namespace nnvm; - using namespace imperative; - - std::vector inputs(cinputs.size()); - std::vector outputs(coutputs.size()); - for (size_t i = 0; i < inputs.size(); i++) - inputs[i] = &cinputs[i]; - for (size_t i = 0; i < outputs.size(); i++) - outputs[i] = &coutputs[i]; - Imperative::CachedOp op(sym, std::vector >()); - op.Forward(nullptr, inputs, outputs); -} - struct ForeachParam : public dmlc::Parameter { int num_args; int dim; @@ -89,19 +73,115 @@ static std::vector ReorderInputs(const std::vector &in, const nnvm::Indexe struct ForeachState { Symbol subgraph; ForeachParam params; + // These are output arrays from all iterations. + // They also contain the Op state for each CachedOp. + std::vector > all_outputs; + std::vector > all_inputs; + std::vector > all_gradients; + std::vector iter_ops; ForeachState(const Symbol &g, const ForeachParam ¶ms) { this->subgraph = g; this->params = params; } + + void Forward(std::vector cinputs, + const std::vector& req, + std::vector coutputs, bool is_recording); + void Backward(int iter_no, std::vector ograds, + const std::vector &req, + std::vector igrads); }; +void ForeachState::Forward(std::vector cinputs, + const std::vector& req, + std::vector coutputs, bool is_recording) { + using namespace nnvm; + using namespace imperative; + + bool orig_is_record; + if (is_recording) + orig_is_record = Imperative::Get()->set_is_recording(true); + else + orig_is_record = Imperative::Get()->is_recording(); + + std::vector inputs(cinputs.size()); + std::vector outputs(coutputs.size()); + for (size_t i = 0; i < inputs.size(); i++) + inputs[i] = &cinputs[i]; + for (size_t i = 0; i < outputs.size(); i++) + outputs[i] = &coutputs[i]; + + if (is_recording) { + all_inputs.push_back(cinputs); + std::vector gradients(cinputs.size()); + std::vector input_ptrs(cinputs.size()); + std::vector gradient_ptrs(cinputs.size()); + std::vector grad_reqs(cinputs.size()); + for (size_t i = 0; i < gradients.size(); i++) { + gradients[i] = NDArray(cinputs[i].shape(), cinputs[i].ctx(), + true, cinputs[i].dtype()); + input_ptrs[i] = &cinputs[i]; + gradient_ptrs[i] = &gradients[i]; + grad_reqs[i] = kWriteTo; + } + Imperative::Get()->MarkVariables(input_ptrs, grad_reqs, gradient_ptrs);; + } + + std::vector > kwargs; + kwargs.push_back(std::pair("inline_limit", "0")); + CachedOpPtr op = std::make_shared(subgraph, kwargs); + // TODO here we only changed the output arrays in the arguments. + // Will this be a problem? + op->Forward(nullptr, inputs, outputs); + + if (is_recording) { + // TODO does this have right inputs and outputs? + all_outputs.push_back(coutputs); + iter_ops.push_back(op); + } + + Imperative::Get()->set_is_recording(orig_is_record); +} + +void ForeachState::Backward(int iter_no, std::vector ograds, + const std::vector &req, + std::vector igrads) { + using namespace nnvm; + using namespace imperative; + + auto op = iter_ops[iter_no]; + std::vector inputs; + std::vector outputs; + inputs.reserve(op->num_backward_inputs()); + outputs.reserve(op->num_inputs()); + for (size_t i = 0; i < ograds.size(); i++) + inputs.push_back(&ograds[i]); +// for (size_t i = 0; i < all_inputs[iter_no].size(); i++) +// inputs.push_back(&all_inputs[iter_no][i]); +// for (size_t i = 0; i < all_outputs[iter_no].size(); i++) +// inputs.push_back(&all_outputs[iter_no][i]); + CHECK_EQ(inputs.size(), op->num_backward_inputs()); + for (size_t i = 0; i < igrads.size(); i++) + outputs.push_back(&igrads[i]); + CHECK_EQ(outputs.size(), op->num_inputs()); + + // TODO here we only changed the output arrays in the arguments. + // Will this be a problem? + CHECK(!Imperative::AGInfo::IsNone(all_outputs[iter_no][0])); + const nnvm::NodeEntry &node_entry = all_outputs[iter_no][0].GetAutogradEntry(); + OpStatePtr state = Imperative::AGInfo::Get(node_entry.node).state; + op->Backward(false, state, inputs, req, outputs); +} + +static bool is_recording = true; + static void ForeachComputeExCPU(const OpStatePtr& state_ptr, const OpContext& ctx, const std::vector& inputs, const std::vector& req, const std::vector& outputs) { - ForeachState state = state_ptr.get_state(); + ForeachState &state = state_ptr.get_state(); const ForeachParam& params = state.params; CHECK_EQ(outputs.size(), (size_t) params.num_outputs); size_t len = inputs[0].shape()[0]; @@ -127,13 +207,13 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, if (len % 2 == 1) { for (size_t i = 1; i < subg_outputs1.size(); i++) { subg_outputs1[i] = outputs[i]; - subg_outputs2[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), false, + subg_outputs2[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), true, outputs[i].dtype()); } } else { // Otherwise, we'll use the second set of outputs. for (size_t i = 1; i < subg_outputs1.size(); i++) { - subg_outputs1[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), false, + subg_outputs1[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), true, outputs[i].dtype()); subg_outputs2[i] = outputs[i]; } @@ -143,9 +223,24 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, for (size_t i = 0; i < len; i++) { std::vector *subg_out_curr = subg_outputs[i % 2]; std::vector *subg_out_prev = subg_outputs[(i + 1) % 2]; + // TODO it might be possible that the data won't be written to the output + // array directly. (*subg_out_curr)[0] = outputs[0].At(i); + // When recording for backward computation, we should make sure + // that output arrays are actually different in each iteration. + if (is_recording && i < len - 1) { + for (size_t j = 1; j < subg_out_curr->size(); j++) + (*subg_out_curr)[j] = NDArray(outputs[j].shape(), outputs[j].ctx(), + true, outputs[j].dtype()); + } else if (is_recording && i == len - 1) { + // For the last iteration, we need to write data to the output array + // directly. + for (size_t j = 1; j < subg_out_curr->size(); j++) + (*subg_out_curr)[j] = outputs[j]; + } // Get a slice from the first input array. + // TODO how can we be sure that the first subgraph input is the data input? subg_inputs[0] = inputs[0].At(i); // For the rest of the iterations, the rest of the arguments are the outputs // from the previous iteration. @@ -156,7 +251,7 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, } } - ExecSubgraph(state.subgraph, ctx, subg_inputs, req, *subg_out_curr); + state.Forward(subg_inputs, req, *subg_out_curr, is_recording); // We need to wait for the iteration to complete before executing // the next one or return from the loop. In this way, we can reuse // the memory in the subgraph. @@ -165,6 +260,51 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, } } +static void ForeachGradComputeExCPU(const OpStatePtr& state_ptr, + const OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs) { + ForeachState &state = state_ptr.get_state(); + const ForeachParam& params = state.params; + CHECK_EQ(outputs.size(), (size_t) params.num_args - 1); + // The inputs contain out gradients, inputs and outputs. + size_t len = inputs[0].shape()[0]; + size_t num_input_data = 1; + size_t num_output_data = 1; + + // In backward computation, we need to run iterations from backwards. + std::vector ograds(params.num_outputs); + std::vector igrads(params.num_args - 1); + for (size_t i = num_output_data; i < ograds.size(); i++) + ograds[i] = inputs[i]; + for (int iter_num = len - 1; iter_num >= 0; iter_num--) { + ograds[0] = inputs[0].At(iter_num); + igrads[0] = outputs[0].At(iter_num); + if (iter_num == 0) { + for (size_t i = num_input_data; i < igrads.size(); i++) + igrads[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), + true, outputs[i].dtype()); + } else { + for (size_t i = num_input_data; i < igrads.size(); i++) + igrads[i] = outputs[i]; + } + + // TODO is req correct here? + state.Backward(iter_num, ograds, req, igrads); + + // We need to wait for the iteration to complete before executing + // the next one or return from the loop. In this way, we can reuse + // the memory in the subgraph. + for (size_t i = 0; i < igrads.size(); i++) + igrads[i].WaitToRead(); + + size_t num_states = ograds.size() - num_output_data; + for (size_t i = 0; i < num_states; i++) + ograds[i + num_output_data] = igrads[i + num_input_data]; + } +} + static bool ForeachShape(const nnvm::NodeAttrs& attrs, std::vector *in_shape, std::vector *out_shape) { @@ -281,14 +421,30 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, return true; } -OpStatePtr CreateForeachState(const NodeAttrs& attrs, - Context ctx, - const std::vector& ishape, - const std::vector& itype) { +static bool BackwardForeachStorageType(const nnvm::NodeAttrs& attrs, + const int dev_mask, + DispatchMode* dispatch_mode, + std::vector *in_attrs, + std::vector *out_attrs) { + // TODO I need to set storage type properly. + return storage_type_assign(out_attrs, mxnet::kDefaultStorage, + dispatch_mode, DispatchMode::kFComputeEx); +} + +static OpStatePtr CreateForeachState(const NodeAttrs& attrs, + Context ctx, + const std::vector& ishape, + const std::vector& itype) { const ForeachParam& params = nnvm::get(attrs.parsed); return OpStatePtr::Create(*attrs.subgraphs[0], params); } +void ForeachParamParser(nnvm::NodeAttrs* attrs) { + ParamParser(attrs); + // This is to indicate that the operator has a subgraph. + attrs->subgraphs.resize(1); +} + NNVM_REGISTER_OP(_foreach) .describe(R"code(foreach)code" ADD_FILELINE) .set_attr_parser(ParamParser) @@ -309,6 +465,7 @@ NNVM_REGISTER_OP(_foreach) [](const NodeAttrs& attrs) { return std::vector{0}; }) +.set_attr("FGradient", ElemwiseGradUseInOut{"_backward_foreach"}) .set_attr("FCreateOpState", CreateForeachState) .set_attr("FInferShape", ForeachShape) .set_attr("FInferType", ForeachType) @@ -319,5 +476,20 @@ NNVM_REGISTER_OP(_foreach) .add_argument("states", "NDArray-or-Symbol[]", "The list of initial states.") .add_arguments(ForeachParam::__FIELDS__()); +NNVM_REGISTER_OP(_backward_foreach) +.set_num_inputs([](const NodeAttrs& attrs){ + const ForeachParam& params = nnvm::get(attrs.parsed); + return params.num_outputs * 2 + params.num_args - 1; + }) +.set_num_outputs([](const NodeAttrs& attrs){ + const ForeachParam& params = nnvm::get(attrs.parsed); + return params.num_args - 1; + }) +.set_attr("FInferStorageType", BackwardForeachStorageType) +.set_attr_parser(ForeachParamParser) +.set_attr("TIsLayerOpBackward", true) +.set_attr("TIsBackward", true) +.set_attr("FStatefulComputeEx", ForeachGradComputeExCPU); + } // namespace op } // namespace mxnet diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 38507771407a..022f0eb526d1 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5678,18 +5678,35 @@ def step(in1, states): out = mx.sym.Group([out1, out[1][0]]) arr1 = mx.nd.random.uniform(shape=(5, 2)) arr2 = mx.nd.random.uniform(shape=(2)) - e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}) - e.forward() - arr1 = arr1.asnumpy() - arr2 = arr2.asnumpy() - np_res = np.zeros_like(arr1) - for i in range(arr1.shape[0]): - if (i == 0): - np_res[i] = arr2 + arr1[i] - else: - np_res[i] = np_res[i - 1] + arr1[i] - np_res = np_res * 2 - assert_almost_equal(e.outputs[0].asnumpy(), np_res, rtol=0.001, atol=0.0001) + arr_grad1 = mx.nd.empty(arr1.shape) + arr_grad2 = mx.nd.empty(arr2.shape) + e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}, + args_grad={'v3': arr_grad1, 'v4': arr_grad2}) + e.forward(is_train=True) + + out_grad = mx.nd.random.uniform(-10, 10, arr1.shape) + state_grad = mx.nd.random.uniform(-10, 10, arr2.shape) + # backward + e.backward([out_grad, state_grad]) + #e.backward() + + res = [] + arr1.attach_grad() + arr2.attach_grad() + with mx.autograd.record(): + for i in range(arr1.shape[0]): + if (i == 0): + tmp_res = mx.nd.expand_dims(arr2, 0) + mx.nd.expand_dims(arr1[i], 0) + else: + tmp_res = res[len(res) - 1] + mx.nd.expand_dims(arr1[i], 0) + res.append(tmp_res) + res1 = mx.nd.concat(*res, dim=0) + res2 = res1 * 2 + res = mx.nd.concat(res2, tmp_res, dim=0) + res.backward(mx.nd.concat(out_grad, mx.nd.expand_dims(state_grad, 0), dim=0)) + assert_almost_equal(e.outputs[0].asnumpy(), res2.asnumpy(), rtol=0.001, atol=0.0001) + assert_almost_equal(arr1.grad.asnumpy(), e.grad_arrays[0].asnumpy()) + assert_almost_equal(arr2.grad.asnumpy(), e.grad_arrays[1].asnumpy()) @with_seed() From 0f111ffc7c1a3d2f9fb3c886fa4468d616b5e145 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 4 May 2018 16:43:36 +0000 Subject: [PATCH 31/71] Fix bugs. --- src/operator/nn/control_flow.cc | 37 ++++++++++++++++++++++++++------- 1 file changed, 30 insertions(+), 7 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index e84b35267cb8..cb5aca4e9ac6 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -91,6 +91,12 @@ struct ForeachState { void Backward(int iter_no, std::vector ograds, const std::vector &req, std::vector igrads); + void Cleanup() { + all_outputs.clear(); + all_inputs.clear(); + all_gradients.clear(); + iter_ops.clear(); + } }; void ForeachState::Forward(std::vector cinputs, @@ -157,10 +163,19 @@ void ForeachState::Backward(int iter_no, std::vector ograds, outputs.reserve(op->num_inputs()); for (size_t i = 0; i < ograds.size(); i++) inputs.push_back(&ograds[i]); -// for (size_t i = 0; i < all_inputs[iter_no].size(); i++) -// inputs.push_back(&all_inputs[iter_no][i]); -// for (size_t i = 0; i < all_outputs[iter_no].size(); i++) -// inputs.push_back(&all_outputs[iter_no][i]); + + const std::vector &save_inputs = op->save_inputs(); + const std::vector &save_outputs = op->save_outputs(); + CHECK_EQ(save_inputs.size(), all_inputs[iter_no].size()); + CHECK_EQ(op->num_outputs(), all_outputs[iter_no].size()); + for (size_t i = 0; i < all_inputs[iter_no].size(); i++) { + if (save_inputs[i]) + inputs.push_back(&all_inputs[iter_no][i]); + } + for (size_t i = 0; i < all_outputs[iter_no].size(); i++) { + if (save_outputs[i]) + inputs.push_back(&all_outputs[iter_no][i]); + } CHECK_EQ(inputs.size(), op->num_backward_inputs()); for (size_t i = 0; i < igrads.size(); i++) outputs.push_back(&igrads[i]); @@ -281,7 +296,11 @@ static void ForeachGradComputeExCPU(const OpStatePtr& state_ptr, for (int iter_num = len - 1; iter_num >= 0; iter_num--) { ograds[0] = inputs[0].At(iter_num); igrads[0] = outputs[0].At(iter_num); - if (iter_num == 0) { + // There are three types of arrays in igrads. + // * data gradients. + // * loop variable gradients. + // * read-only variable gradients. + if (iter_num != 0) { for (size_t i = num_input_data; i < igrads.size(); i++) igrads[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), true, outputs[i].dtype()); @@ -300,9 +319,13 @@ static void ForeachGradComputeExCPU(const OpStatePtr& state_ptr, igrads[i].WaitToRead(); size_t num_states = ograds.size() - num_output_data; - for (size_t i = 0; i < num_states; i++) - ograds[i + num_output_data] = igrads[i + num_input_data]; + for (size_t i = 0; i < num_states; i++) { + size_t loc = params.in_state_locs[i]; + CHECK_LT(loc, igrads.size()); + ograds[i + num_output_data] = igrads[loc]; + } } + state.Cleanup(); } static bool ForeachShape(const nnvm::NodeAttrs& attrs, From 7688b740ac28d364a18ec064e652f2b49530ec09 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 4 May 2018 16:44:28 +0000 Subject: [PATCH 32/71] enable backward test in lstm. --- tests/python/unittest/test_operator.py | 55 ++++++++++++++++++++------ 1 file changed, 42 insertions(+), 13 deletions(-) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 022f0eb526d1..2834fc2eac76 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5683,12 +5683,10 @@ def step(in1, states): e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}, args_grad={'v3': arr_grad1, 'v4': arr_grad2}) e.forward(is_train=True) - + # backward out_grad = mx.nd.random.uniform(-10, 10, arr1.shape) state_grad = mx.nd.random.uniform(-10, 10, arr2.shape) - # backward e.backward([out_grad, state_grad]) - #e.backward() res = [] arr1.attach_grad() @@ -5747,13 +5745,37 @@ def sym_group(out): i2h_barr = mx.nd.random.uniform(shape=(16)) h2h_barr = mx.nd.random.uniform(shape=(16)) + data_arr_grad1 = mx.nd.empty(data_arr.shape) + h_arr_grad1 = mx.nd.empty(h_arr.shape) + c_arr_grad1 = mx.nd.empty(c_arr.shape) + i2h_warr_grad1 = mx.nd.empty(i2h_warr.shape) + h2h_warr_grad1 = mx.nd.empty(h2h_warr.shape) + i2h_barr_grad1 = mx.nd.empty(i2h_barr.shape) + h2h_barr_grad1 = mx.nd.empty(h2h_barr.shape) out = mx.sym.contrib.foreach(step, data, [init_h, init_c]) out = sym_group(out) - e = out.bind(ctx=mx.cpu(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, - 'i2h_weight': i2h_warr, 'h2h_weight': h2h_warr, 'i2h_bias': i2h_barr, 'h2h_bias': h2h_barr}) - e.forward() - outputs1 = e.outputs - + e1 = out.bind(ctx=mx.cpu(), + args={'data': data_arr, 'h': h_arr, 'c': c_arr, + 'i2h_weight': i2h_warr, 'h2h_weight': h2h_warr, + 'i2h_bias': i2h_barr, 'h2h_bias': h2h_barr}, + args_grad={'data': data_arr_grad1, 'h': h_arr_grad1, 'c': c_arr_grad1, + 'i2h_weight': i2h_warr_grad1, 'h2h_weight': h2h_warr_grad1, + 'i2h_bias': i2h_barr_grad1, 'h2h_bias': h2h_barr_grad1}) + e1.forward(is_train=True) + outputs1 = e1.outputs + # backward + out_grads = [] + for arr in e1.outputs: + out_grads.append(mx.nd.random.uniform(-10, 10, arr.shape)) + e1.backward(out_grads) + + data_arr_grad2 = mx.nd.empty(data_arr.shape) + h_arr_grad2 = mx.nd.empty(h_arr.shape) + c_arr_grad2 = mx.nd.empty(c_arr.shape) + i2h_warr_grad2 = mx.nd.empty(i2h_warr.shape) + h2h_warr_grad2 = mx.nd.empty(h2h_warr.shape) + i2h_barr_grad2 = mx.nd.empty(i2h_barr.shape) + h2h_barr_grad2 = mx.nd.empty(h2h_barr.shape) lstm = mx.rnn.LSTMCell(4, prefix='mylstm_') h = init_h c = init_c @@ -5763,14 +5785,21 @@ def sym_group(out): unroll_outs.append(mx.sym.expand_dims(h, axis=0)) unroll_outs = mx.sym.concat(*unroll_outs, dim=0) out = mx.sym.Group([unroll_outs, h, c]) - e = out.bind(ctx=mx.cpu(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, - 'mylstm_i2h_weight': i2h_warr, 'mylstm_h2h_weight': h2h_warr, - 'mylstm_i2h_bias': i2h_barr, 'mylstm_h2h_bias': h2h_barr}) - e.forward() - outputs2 = e.outputs + e2 = out.bind(ctx=mx.cpu(), + args={'data': data_arr, 'h': h_arr, 'c': c_arr, + 'mylstm_i2h_weight': i2h_warr, 'mylstm_h2h_weight': h2h_warr, + 'mylstm_i2h_bias': i2h_barr, 'mylstm_h2h_bias': h2h_barr}, + args_grad={'data': data_arr_grad2, 'h': h_arr_grad2, 'c': c_arr_grad2, + 'mylstm_i2h_weight': i2h_warr_grad2, 'mylstm_h2h_weight': h2h_warr_grad2, + 'mylstm_i2h_bias': i2h_barr_grad2, 'mylstm_h2h_bias': h2h_barr_grad2}) + e2.forward(is_train=True) + outputs2 = e2.outputs + e2.backward(out_grads) for i in range(len(outputs2)): assert_almost_equal(outputs1[i].asnumpy(), outputs2[i].asnumpy(), rtol=0.001, atol=0.0001) + for i in range(len(e1.grad_arrays)): + assert_almost_equal(e1.grad_arrays[i].asnumpy(), e2.grad_arrays[i].asnumpy()) @with_seed() From f1141a43ff964a1e9104b2fad23b7324d6303ad2 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 7 May 2018 17:28:46 +0000 Subject: [PATCH 33/71] Fix a bug in foreach backward for free variables. --- src/operator/nn/control_flow.cc | 45 ++++++++++++++++++-------- tests/python/unittest/test_operator.py | 25 ++++++++++---- 2 files changed, 50 insertions(+), 20 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index cb5aca4e9ac6..4a31af53a290 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -284,39 +284,58 @@ static void ForeachGradComputeExCPU(const OpStatePtr& state_ptr, const ForeachParam& params = state.params; CHECK_EQ(outputs.size(), (size_t) params.num_args - 1); // The inputs contain out gradients, inputs and outputs. - size_t len = inputs[0].shape()[0]; + int len = inputs[0].shape()[0]; size_t num_input_data = 1; size_t num_output_data = 1; // In backward computation, we need to run iterations from backwards. std::vector ograds(params.num_outputs); - std::vector igrads(params.num_args - 1); + std::vector igrads(outputs.size()); for (size_t i = num_output_data; i < ograds.size(); i++) ograds[i] = inputs[i]; + std::vector iter_req(req.size()); + for (auto r : req) + CHECK_NE(r, kWriteInplace); for (int iter_num = len - 1; iter_num >= 0; iter_num--) { + // TODO data isn't always the first one. ograds[0] = inputs[0].At(iter_num); igrads[0] = outputs[0].At(iter_num); - // There are three types of arrays in igrads. - // * data gradients. - // * loop variable gradients. - // * read-only variable gradients. - if (iter_num != 0) { - for (size_t i = num_input_data; i < igrads.size(); i++) + iter_req[0] = req[0]; + for (size_t i = num_input_data; i < igrads.size(); i++) { + // There are three types of arrays in igrads. + // * data gradients. + // * loop variable gradients. + // * read-only variable gradients. + // For state gradients, we need to allocate new NDArrays + // because intermediate state gradients won't be returned to the users. + bool in_state = std::find(params.in_state_locs.begin(), + params.in_state_locs.end(), + i) != params.in_state_locs.end(); + if (iter_num != 0 && in_state) { igrads[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), true, outputs[i].dtype()); - } else { - for (size_t i = num_input_data; i < igrads.size(); i++) + } else { igrads[i] = outputs[i]; + } + if (in_state) + // For the first iteration, we need to use the request provided by + // the user to write state gradients to the outputs. + iter_req[i] = iter_num != 0 ? kWriteTo : req[i]; + else + // For all read-only variable gradients, we need to use the request + // provided by the user in the last iteration and later on add gradients + // to the output arrays. + iter_req[i] = iter_num == len - 1 ? req[i]: kAddTo; } - // TODO is req correct here? - state.Backward(iter_num, ograds, req, igrads); + state.Backward(iter_num, ograds, iter_req, igrads); // We need to wait for the iteration to complete before executing // the next one or return from the loop. In this way, we can reuse // the memory in the subgraph. - for (size_t i = 0; i < igrads.size(); i++) + for (size_t i = 0; i < igrads.size(); i++) { igrads[i].WaitToRead(); + } size_t num_states = ograds.size() - num_output_data; for (size_t i = 0; i < num_states; i++) { diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 2834fc2eac76..6124fe093770 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5667,21 +5667,24 @@ def test_float16_min_max(): def test_foreach(): v3 = mx.sym.var("v3") v4 = mx.sym.var("v4") + v5 = mx.sym.var("v5") # This tests foreach with accumulation sum. def step(in1, states): - out = in1 + states[0] + out = in1 + states[0] + v5 return (out, [out]) out = mx.sym.contrib.foreach(step, v3, [v4]) out1 = out[0] * 2 out = mx.sym.Group([out1, out[1][0]]) - arr1 = mx.nd.random.uniform(shape=(5, 2)) + arr1 = mx.nd.random.uniform(shape=(2, 2)) arr2 = mx.nd.random.uniform(shape=(2)) + arr3 = mx.nd.random.uniform(shape=(2)) arr_grad1 = mx.nd.empty(arr1.shape) arr_grad2 = mx.nd.empty(arr2.shape) - e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2}, - args_grad={'v3': arr_grad1, 'v4': arr_grad2}) + arr_grad3 = mx.nd.empty(arr3.shape) + e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2, 'v5': arr3}, + args_grad={'v3': arr_grad1, 'v4': arr_grad2, 'v5': arr_grad3}) e.forward(is_train=True) # backward out_grad = mx.nd.random.uniform(-10, 10, arr1.shape) @@ -5691,12 +5694,13 @@ def step(in1, states): res = [] arr1.attach_grad() arr2.attach_grad() + arr3.attach_grad() with mx.autograd.record(): for i in range(arr1.shape[0]): if (i == 0): - tmp_res = mx.nd.expand_dims(arr2, 0) + mx.nd.expand_dims(arr1[i], 0) + tmp_res = mx.nd.expand_dims(arr2, 0) + mx.nd.expand_dims(arr1[i], 0) + mx.nd.expand_dims(arr3, 0) else: - tmp_res = res[len(res) - 1] + mx.nd.expand_dims(arr1[i], 0) + tmp_res = res[len(res) - 1] + mx.nd.expand_dims(arr1[i], 0) + mx.nd.expand_dims(arr3, 0) res.append(tmp_res) res1 = mx.nd.concat(*res, dim=0) res2 = res1 * 2 @@ -5705,6 +5709,7 @@ def step(in1, states): assert_almost_equal(e.outputs[0].asnumpy(), res2.asnumpy(), rtol=0.001, atol=0.0001) assert_almost_equal(arr1.grad.asnumpy(), e.grad_arrays[0].asnumpy()) assert_almost_equal(arr2.grad.asnumpy(), e.grad_arrays[1].asnumpy()) + assert_almost_equal(arr3.grad.asnumpy(), e.grad_arrays[2].asnumpy()) @with_seed() @@ -5737,7 +5742,7 @@ def sym_group(out): ret.extend(out[1]) return mx.sym.Group(ret) - data_arr = mx.nd.random.uniform(shape=(5, 2, 4)) + data_arr = mx.nd.random.uniform(shape=(2, 2, 4)) h_arr = mx.nd.random.uniform(shape=(2, 4)) c_arr = mx.nd.random.uniform(shape=(2, 4)) i2h_warr = mx.nd.random.uniform(shape=(16, 4)) @@ -5802,6 +5807,12 @@ def sym_group(out): assert_almost_equal(e1.grad_arrays[i].asnumpy(), e2.grad_arrays[i].asnumpy()) +# TODO Test cases: +# in an iteration, data is stored in any location. +# # iterations: odd or even. +# multiple inputs and multiple outputs. + + @with_seed() def test_squeeze_op(): def check_squeeze_op(shape, axis=None): From 8f8e51d0f5cb572ba943f07d3333e45d6c59e253 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 9 May 2018 00:31:33 +0000 Subject: [PATCH 34/71] change for the new CachedOp. --- src/operator/nn/control_flow.cc | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 4a31af53a290..8d5ede3391b9 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -71,7 +71,8 @@ static std::vector ReorderInputs(const std::vector &in, const nnvm::Indexe } struct ForeachState { - Symbol subgraph; + Symbol subgraph_sym; + nnvm::Graph subgraph; ForeachParam params; // These are output arrays from all iterations. // They also contain the Op state for each CachedOp. @@ -81,7 +82,8 @@ struct ForeachState { std::vector iter_ops; ForeachState(const Symbol &g, const ForeachParam ¶ms) { - this->subgraph = g; + this->subgraph_sym = g; + this->subgraph.outputs = g.outputs; this->params = params; } @@ -136,7 +138,15 @@ void ForeachState::Forward(std::vector cinputs, std::vector > kwargs; kwargs.push_back(std::pair("inline_limit", "0")); - CachedOpPtr op = std::make_shared(subgraph, kwargs); + // Get input names. + const auto& idx = subgraph.indexed_graph(); + std::vector arg_names(idx.input_nodes().size()); + for (size_t i = 0; i < idx.input_nodes().size(); ++i) + arg_names[i] = idx[idx.input_nodes()[i]].source->attrs.name; + // We don't have parameters for the cached op. + std::unordered_map > params; + CachedOpPtr op = std::make_shared(subgraph_sym, kwargs, + arg_names, params); // TODO here we only changed the output arrays in the arguments. // Will this be a problem? op->Forward(nullptr, inputs, outputs); From 05ed08d7849197d7578203d9e5abf083df0bce3d Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 9 May 2018 03:33:08 +0000 Subject: [PATCH 35/71] Detect the backward computation. --- include/mxnet/op_attr_types.h | 4 +++- src/executor/graph_executor.cc | 9 ++++++--- src/executor/graph_executor.h | 2 ++ src/imperative/imperative_utils.h | 12 ++++++++---- src/ndarray/ndarray.cc | 3 ++- src/operator/nn/control_flow.cc | 20 +++++++++++--------- 6 files changed, 32 insertions(+), 18 deletions(-) diff --git a/include/mxnet/op_attr_types.h b/include/mxnet/op_attr_types.h index 3969d8445be1..23a318464f15 100644 --- a/include/mxnet/op_attr_types.h +++ b/include/mxnet/op_attr_types.h @@ -64,8 +64,10 @@ enum OpReqType { * \sa Resource */ struct OpContext { + /*! \brief whether there is a backward phase to compute gradients. */ + bool need_grad; /*! \brief whether it is training phase */ - int is_train; + bool is_train; /*! \brief RunContext related resources */ RunContext run_ctx; /*! \brief the callback when operation completes, used by asynchronize ops */ diff --git a/src/executor/graph_executor.cc b/src/executor/graph_executor.cc index fc668b733639..cf1263b79057 100644 --- a/src/executor/graph_executor.cc +++ b/src/executor/graph_executor.cc @@ -39,6 +39,7 @@ namespace exec { GraphExecutor::GraphExecutor() { log_verbose_ = dmlc::GetEnv("MXNET_EXEC_VERBOSE_LOGGING", false); + need_grad_ = false; } GraphExecutor::~GraphExecutor() { @@ -257,11 +258,11 @@ nnvm::Graph GraphExecutor::InitFullGraph(nnvm::Symbol symbol, nnvm::Graph g; g.outputs = symbol.outputs; - bool need_grad = false; + need_grad_ = false; for (OpReqType req : grad_req_types) { - if (req != kNullOp) need_grad = true; + if (req != kNullOp) need_grad_ = true; } - if (!need_grad) return g; + if (!need_grad_) return g; for (size_t i = 0; i < g.outputs.size(); ++i) { NodeEntry ngrad{nnvm::Node::Create(), 0, 0}; head_grad_entry_.emplace_back(AttrHint(ngrad, g.outputs[i])); @@ -1492,6 +1493,7 @@ void GraphExecutor::RunOps(bool is_train, size_t topo_start, size_t topo_end) { const auto& inode = idx[nid]; if (inode.source->is_variable()) continue; opnode.exec->op_ctx.is_train = is_train; + opnode.exec->op_ctx.need_grad = need_grad_; } // Push Ops @@ -1510,6 +1512,7 @@ void GraphExecutor::RunOps(bool is_train, size_t topo_start, size_t topo_end) { OpNode& opnode = op_nodes_[nid]; if (op_nodes_[nid].skip_exec_node) continue; opnode.exec->op_ctx.is_train = is_train; + opnode.exec->op_ctx.need_grad = need_grad_; if (opnode.exec->exec_type() == ExecType::kCrossDeviceCopy) { CHECK_EQ(inode.inputs.size(), 1U); CHECK_EQ(opnode.exec->in_array.size(), 1U); diff --git a/src/executor/graph_executor.h b/src/executor/graph_executor.h index bcde41d508eb..fa2a156d3d76 100644 --- a/src/executor/graph_executor.h +++ b/src/executor/graph_executor.h @@ -203,6 +203,8 @@ class GraphExecutor : public Executor { // perform bulking and segmentation on a training graph void BulkTrainingOpSegs(size_t total_num_nodes); + // indicate whether there is a backward graph for gradients. + bool need_grad_; // internal graph nnvm::Graph graph_; // operator node diff --git a/src/imperative/imperative_utils.h b/src/imperative/imperative_utils.h index d7bb37b7cfef..35c8bd4d2bf4 100644 --- a/src/imperative/imperative_utils.h +++ b/src/imperative/imperative_utils.h @@ -379,7 +379,8 @@ inline void PushFCompute(const FCompute& fn, &input_blobs, &output_blobs, &pre_temp_src, &pre_temp_dst, &post_temp_src, &post_temp_dst, &in_temp_idx_map, mutate_idx); // setup context - OpContext opctx{is_train, rctx, engine::CallbackOnComplete(), requested}; + bool need_grad = Imperative::Get()->is_recording(); + OpContext opctx{need_grad, is_train, rctx, engine::CallbackOnComplete(), requested}; bool is_gpu = ctx.dev_mask() == gpu::kDevMask; // pre-fcompute fallback, cast to default storage type CastNonDefaultStorage(pre_temp_src, pre_temp_dst, opctx, is_gpu); @@ -410,7 +411,8 @@ inline void PushFComputeEx(const FComputeEx& fn, std::vector inputs, outputs; DerefInputOutput(p_inputs, p_outputs, &inputs, &outputs); const auto& run = [=](RunContext rctx) { - OpContext opctx{is_train, rctx, engine::CallbackOnComplete(), requested}; + bool need_grad = Imperative::Get()->is_recording(); + OpContext opctx{need_grad, is_train, rctx, engine::CallbackOnComplete(), requested}; #if MXNET_USE_MKLDNN == 1 InvalidateOutputs(outputs, req); #endif @@ -456,7 +458,8 @@ inline void PushOperator(const OpStatePtr& state, if (fcompute_ex != nullptr && dispatch_mode == DispatchMode::kFComputeEx) { const auto& run = [=](RunContext rctx, engine::CallbackOnComplete on_complete) { - OpContext opctx{is_train, rctx, on_complete, requested}; + bool need_grad = Imperative::Get()->is_recording(); + OpContext opctx{need_grad, is_train, rctx, on_complete, requested}; #if MXNET_USE_MKLDNN == 1 InvalidateOutputs(outputs, req); #endif @@ -483,7 +486,8 @@ inline void PushOperator(const OpStatePtr& state, << "for stateful operator " << op->name; const auto& run = [=](RunContext rctx, engine::CallbackOnComplete on_complete) { - OpContext opctx{is_train, rctx, on_complete, requested}; + bool need_grad = Imperative::Get()->is_recording(); + OpContext opctx{need_grad, is_train, rctx, on_complete, requested}; std::vector input_blobs, output_blobs; // pre-fcompute and post-fcompute storage fallback src NDArrays and dst NDArrays diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc index d87e8bc95ea5..1accd6f9f1cb 100644 --- a/src/ndarray/ndarray.cc +++ b/src/ndarray/ndarray.cc @@ -1160,7 +1160,8 @@ void CopyFromToImpl(const NDArray& from, const NDArray& to, const Context to_ctx = to.ctx(); bool is_train = Imperative::Get()->is_training(); - OpContext opctx{is_train, + OpContext opctx{Imperative::Get()->is_recording(), + is_train, rctx, engine::CallbackOnComplete(), requested}; diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 8d5ede3391b9..0a3d2e1a78b4 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -70,10 +70,7 @@ static std::vector ReorderInputs(const std::vector &in, const nnvm::Indexe return ret; } -struct ForeachState { - Symbol subgraph_sym; - nnvm::Graph subgraph; - ForeachParam params; +class ForeachState { // These are output arrays from all iterations. // They also contain the Op state for each CachedOp. std::vector > all_outputs; @@ -81,6 +78,11 @@ struct ForeachState { std::vector > all_gradients; std::vector iter_ops; + public: + Symbol subgraph_sym; + nnvm::Graph subgraph; + ForeachParam params; + ForeachState(const Symbol &g, const ForeachParam ¶ms) { this->subgraph_sym = g; this->subgraph.outputs = g.outputs; @@ -166,6 +168,8 @@ void ForeachState::Backward(int iter_no, std::vector ograds, using namespace nnvm; using namespace imperative; + CHECK_GT(iter_ops.size(), iter_no) + << "We didn't record the computation for iteration " << iter_no; auto op = iter_ops[iter_no]; std::vector inputs; std::vector outputs; @@ -199,8 +203,6 @@ void ForeachState::Backward(int iter_no, std::vector ograds, op->Backward(false, state, inputs, req, outputs); } -static bool is_recording = true; - static void ForeachComputeExCPU(const OpStatePtr& state_ptr, const OpContext& ctx, const std::vector& inputs, @@ -253,11 +255,11 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, (*subg_out_curr)[0] = outputs[0].At(i); // When recording for backward computation, we should make sure // that output arrays are actually different in each iteration. - if (is_recording && i < len - 1) { + if (ctx.need_grad && i < len - 1) { for (size_t j = 1; j < subg_out_curr->size(); j++) (*subg_out_curr)[j] = NDArray(outputs[j].shape(), outputs[j].ctx(), true, outputs[j].dtype()); - } else if (is_recording && i == len - 1) { + } else if (ctx.need_grad && i == len - 1) { // For the last iteration, we need to write data to the output array // directly. for (size_t j = 1; j < subg_out_curr->size(); j++) @@ -276,7 +278,7 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, } } - state.Forward(subg_inputs, req, *subg_out_curr, is_recording); + state.Forward(subg_inputs, req, *subg_out_curr, ctx.need_grad); // We need to wait for the iteration to complete before executing // the next one or return from the loop. In this way, we can reuse // the memory in the subgraph. From 610a9dc4d4dcd364d9e5ace4a2ce2948aec6bc92 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 9 May 2018 23:26:34 +0000 Subject: [PATCH 36/71] Fix bugs in foreach. --- python/mxnet/symbol/contrib.py | 32 ++++--- src/operator/nn/control_flow.cc | 152 ++++++++++++++++++++------------ 2 files changed, 116 insertions(+), 68 deletions(-) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index 856591ca0c09..1d82cdc92763 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -30,7 +30,7 @@ from . import symbol from ..base import _LIB, c_str, c_array, check_call -from ..base import SymbolHandle, NDArrayHandle +from ..base import SymbolHandle, NDArrayHandle, _as_list from ..attribute import AttrScope __all__ = ["rand_zipfian"] @@ -115,11 +115,14 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): assert isinstance(init_states, list), "init_states should be a list" states = [] with AttrScope(subgraph_name=name): - in_ele = symbol.var("in") + if isinstance(input, list): + in_eles = [symbol.var(sym.name) for sym in input] + else: + in_eles = symbol.var(input.name) for s in init_states: states.append(symbol.var(s.name)) - sym_out = func(in_ele, states) + sym_out = func(in_eles, states) # The function should return a tuple. The first element goes to # the output of the function. The second element is a list. assert isinstance(sym_out, tuple), "func should return a tuple (out, states)" @@ -133,6 +136,7 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): flat_out = sym_out[0] else: flat_out = [sym_out[0]] + num_out_data = len(flat_out) for s in sym_out[1]: # There is a problem if the outputs are the same as the inputs # or the first output. @@ -153,26 +157,32 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): input_syms = {sym.name:sym for sym in input_syms} gin_names = input_syms.keys() # This array contains the symbols for the inputs of foreach. + # They are ordered according to the inputs of the subgraph. ordered_ins = [] states_map = {sym.name:sym for sym in init_states} state_names = states_map.keys() - in_state_locs = [-1] * len(init_states) + data_syms = _as_list(input) + data_map = {sym.name:sym for sym in data_syms} + data_names = data_map.keys() + in_state_locs = [] + in_data_locs = [] for in_name in g.list_inputs(): assert in_name in gin_names, "The input variable %s can't be found in graph inputs: %s" \ % (in_name, str(gin_names)) if (in_name in state_names): ordered_ins.append(states_map[in_name]) - elif (in_name != "in"): + in_state_locs.append(len(ordered_ins) - 1) + elif (in_name in data_names): + ordered_ins.append(data_map[in_name]) + in_data_locs.append(len(ordered_ins) - 1) + else: ordered_ins.append(input_syms[in_name]) - for i in range(len(init_states)): - if (init_states[i].name == in_name): - in_state_locs[i] = len(ordered_ins) - 1 + num_inputs - num_outputs = len(flat_out) num_states = len(state_names) - ret = symbol._internal._foreach(g, input, *ordered_ins, num_outputs=num_outputs, - in_state_locs=in_state_locs) + ret = symbol._internal._foreach(g, *ordered_ins, num_outputs=num_outputs, + num_out_data=num_out_data, in_state_locs=in_state_locs, + in_data_locs=in_data_locs) if (num_outputs - num_states > 1): outs = [] for i in range(num_outputs - num_states): diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 0a3d2e1a78b4..6fe3675dfa5d 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -35,7 +35,9 @@ struct ForeachParam : public dmlc::Parameter { int num_args; int dim; int num_outputs; + int num_out_data; nnvm::Tuple in_state_locs; + nnvm::Tuple in_data_locs; DMLC_DECLARE_PARAMETER(ForeachParam) { DMLC_DECLARE_FIELD(num_args).set_lower_bound(1) .describe("Number of inputs."); @@ -43,8 +45,12 @@ struct ForeachParam : public dmlc::Parameter { .describe("the dimension of the input array to iterate."); DMLC_DECLARE_FIELD(num_outputs) .describe("The number of outputs of the subgraph."); + DMLC_DECLARE_FIELD(num_out_data) + .describe("The number of output data of the subgraph."); DMLC_DECLARE_FIELD(in_state_locs) .describe("The locations of loop states among the inputs."); + DMLC_DECLARE_FIELD(in_data_locs) + .describe("The locations of input data among the inputs."); } }; // struct ForeachParam @@ -210,16 +216,17 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, const std::vector& outputs) { ForeachState &state = state_ptr.get_state(); const ForeachParam& params = state.params; + size_t iter_dim = 0; CHECK_EQ(outputs.size(), (size_t) params.num_outputs); - size_t len = inputs[0].shape()[0]; - CHECK_EQ(inputs[0].shape()[0], outputs[0].shape()[0]); - - // Initialize the inputs for the subgraph. - std::vector subg_inputs(inputs.size()); - for (size_t i = 1; i < inputs.size(); i++) { - // These are the initial states. - subg_inputs[i] = inputs[i]; + CHECK_GT(params.in_data_locs.ndim(), 0); + size_t loc0 = params.in_data_locs[0]; + size_t len = inputs[loc0].shape()[iter_dim]; + for (size_t i = 1; i < params.in_data_locs.ndim(); i++) { + size_t loc = params.in_data_locs[i]; + CHECK_EQ(inputs[loc].shape()[iter_dim], len); } + for (size_t i = 0; i < (size_t) params.num_out_data; i++) + CHECK_EQ(len, outputs[i].shape()[iter_dim]); // Initialize the outputs of the subgraph is a little trickier. // The states from the previous iteration are used as the inputs of the next @@ -246,35 +253,49 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, } } + // Initialize the inputs for the subgraph. + // In each iteration, we need to update the subgraph inputs for input data + // and the loop states. This initialization helps to get the read-only + // arrays in the loop. + std::vector subg_inputs(inputs.size()); + for (size_t i = 0; i < inputs.size(); i++) { + // These are the initial states. + subg_inputs[i] = inputs[i]; + } + // Here we iterate over the first dimension of the first input array. for (size_t i = 0; i < len; i++) { + // Initialize outputs for the subgraph. std::vector *subg_out_curr = subg_outputs[i % 2]; std::vector *subg_out_prev = subg_outputs[(i + 1) % 2]; - // TODO it might be possible that the data won't be written to the output - // array directly. - (*subg_out_curr)[0] = outputs[0].At(i); + for (int j = 0; j < params.num_out_data; j++) + (*subg_out_curr)[j] = outputs[j].At(i); // When recording for backward computation, we should make sure // that output arrays are actually different in each iteration. if (ctx.need_grad && i < len - 1) { - for (size_t j = 1; j < subg_out_curr->size(); j++) + for (size_t j = params.num_out_data; j < subg_out_curr->size(); j++) (*subg_out_curr)[j] = NDArray(outputs[j].shape(), outputs[j].ctx(), true, outputs[j].dtype()); } else if (ctx.need_grad && i == len - 1) { // For the last iteration, we need to write data to the output array // directly. - for (size_t j = 1; j < subg_out_curr->size(); j++) + for (size_t j = params.num_out_data; j < subg_out_curr->size(); j++) (*subg_out_curr)[j] = outputs[j]; } - // Get a slice from the first input array. - // TODO how can we be sure that the first subgraph input is the data input? - subg_inputs[0] = inputs[0].At(i); + // Initialize inputs for the subgraph. + // Get a slice from the input data arrays. + for (size_t j = 0; j < params.in_data_locs.ndim(); j++) { + size_t loc = params.in_data_locs[j]; + subg_inputs[loc] = inputs[loc].At(i); + } // For the rest of the iterations, the rest of the arguments are the outputs // from the previous iteration. if (i > 0) { - for (size_t j = 1; j < subg_out_prev->size(); j++) { - CHECK_LT(params.in_state_locs[j - 1], subg_inputs.size()); - subg_inputs[params.in_state_locs[j - 1]] = (*subg_out_prev)[j]; + for (size_t j = params.num_out_data; j < subg_out_prev->size(); j++) { + size_t idx = j - params.num_out_data; + CHECK_LT(params.in_state_locs[idx], subg_inputs.size()); + subg_inputs[params.in_state_locs[idx]] = (*subg_out_prev)[j]; } } @@ -282,8 +303,9 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, // We need to wait for the iteration to complete before executing // the next one or return from the loop. In this way, we can reuse // the memory in the subgraph. - for (size_t j = 0; j < subg_out_curr->size(); j++) + for (size_t j = 0; j < subg_out_curr->size(); j++) { (*subg_out_curr)[j].WaitToRead(); + } } } @@ -295,10 +317,15 @@ static void ForeachGradComputeExCPU(const OpStatePtr& state_ptr, ForeachState &state = state_ptr.get_state(); const ForeachParam& params = state.params; CHECK_EQ(outputs.size(), (size_t) params.num_args - 1); + CHECK_GT(params.in_data_locs.ndim(), 0); + size_t iter_dim = 0; + std::unordered_set in_data_locs(params.in_data_locs.begin(), + params.in_data_locs.end()); + std::unordered_set in_state_locs(params.in_state_locs.begin(), + params.in_state_locs.end()); // The inputs contain out gradients, inputs and outputs. - int len = inputs[0].shape()[0]; - size_t num_input_data = 1; - size_t num_output_data = 1; + size_t len = inputs[0].shape()[iter_dim]; + size_t num_output_data = params.num_out_data; // In backward computation, we need to run iterations from backwards. std::vector ograds(params.num_outputs); @@ -309,21 +336,26 @@ static void ForeachGradComputeExCPU(const OpStatePtr& state_ptr, for (auto r : req) CHECK_NE(r, kWriteInplace); for (int iter_num = len - 1; iter_num >= 0; iter_num--) { - // TODO data isn't always the first one. - ograds[0] = inputs[0].At(iter_num); - igrads[0] = outputs[0].At(iter_num); - iter_req[0] = req[0]; - for (size_t i = num_input_data; i < igrads.size(); i++) { - // There are three types of arrays in igrads. - // * data gradients. - // * loop variable gradients. - // * read-only variable gradients. - // For state gradients, we need to allocate new NDArrays - // because intermediate state gradients won't be returned to the users. - bool in_state = std::find(params.in_state_locs.begin(), - params.in_state_locs.end(), - i) != params.in_state_locs.end(); + for (int i = 0; i < params.num_out_data; i++) + ograds[i] = inputs[i].At(iter_num); + + // There are three types of arrays in igrads. + // * data gradients. + // * loop variable gradients. + // * read-only variable gradients. + // These are the input data gradients. + for (size_t i = 0; i < igrads.size(); i++) { + // data gradients. + if (in_data_locs.count(i)) { + igrads[i] = outputs[i].At(iter_num); + iter_req[i] = req[i]; + continue; + } + + bool in_state = in_state_locs.count(i); if (iter_num != 0 && in_state) { + // For state gradients, we need to allocate new NDArrays + // because intermediate state gradients won't be returned to the users. igrads[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), true, outputs[i].dtype()); } else { @@ -366,7 +398,13 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, CHECK_EQ(out_shape->size(), (size_t) params.num_outputs); nnvm::ShapeVector shape_inputs = *in_shape; // foreach iterates over the first input NDArray over the first dimension. - shape_inputs[0] = TShape(in_shape->at(0).begin() + 1, in_shape->at(0).end()); + size_t loc0 = params.in_data_locs[0]; + size_t len = in_shape->at(loc0)[0]; + for (size_t i = 0; i < params.in_data_locs.ndim(); i++) { + size_t loc = params.in_data_locs[i]; + CHECK_EQ(len, in_shape->at(loc)[0]); + shape_inputs[loc] = TShape(in_shape->at(loc).begin() + 1, in_shape->at(loc).end()); + } CHECK_EQ(attrs.subgraphs.size(), 1U); auto g = std::make_shared(); g->outputs = attrs.subgraphs[0]->outputs; @@ -381,24 +419,26 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, // We need to copy the inferred input shapes back. const auto &input_nids = idx.input_nodes(); CHECK_EQ(input_nids.size(), in_shape->size()); - size_t num_input_arrays = 1; - for (size_t i = num_input_arrays; i < in_shape->size(); i++) { + for (size_t i = 0; i < in_shape->size(); i++) { auto eid = idx.entry_id(input_nids[i], 0); - (*in_shape)[i] = shapes[eid]; + // If the input shape is none, we should update them. + if ((*in_shape)[i].ndim() == 0 || (*in_shape)[i].Size() == 0) + (*in_shape)[i] = shapes[eid]; } - // For the first shape. - uint32_t eid = idx.entry_id(g->outputs[0]); - const auto& g_out_shape = shapes[eid]; - const auto &in0 = (*in_shape)[0]; - auto &out0 = (*out_shape)[0]; - CHECK_EQ(g_out_shape.ndim() + 1, in0.ndim()); - out0 = in0; - for (size_t i = 1; i < out0.ndim(); i++) - out0[i] = g_out_shape[i - 1]; + // For the shape of output data. + for (int i = 0; i < params.num_out_data; i++) { + uint32_t eid = idx.entry_id(g->outputs[i]); + const auto& g_out_shape = shapes[eid]; + auto &out = (*out_shape)[i]; + out = TShape(g_out_shape.ndim() + 1); + out[0] = len; + for (size_t i = 1; i < out.ndim(); i++) + out[i] = g_out_shape[i - 1]; + } // For the remaining shapes. - for (size_t i = 1; i < g->outputs.size(); i++) { + for (size_t i = params.num_out_data; i < g->outputs.size(); i++) { uint32_t eid = idx.entry_id(g->outputs[i]); (*out_shape)[i] = shapes[eid]; } @@ -419,14 +459,13 @@ static bool ForeachType(const nnvm::NodeAttrs& attrs, CHECK_EQ(idx.outputs().size(), out_type->size()); imperative::CheckAndInferType(g.get(), std::move(dtype_inputs), true); - size_t num_input_arrays = 1; const auto &dtypes = g->GetAttr("dtype"); // Inferring the data type in the subgraph may infer the data type of the inputs. // We need to copy the inferred input data types back. const auto &input_nids = idx.input_nodes(); CHECK_EQ(input_nids.size(), in_type->size()); - for (size_t i = num_input_arrays; i < in_type->size(); i++) { + for (size_t i = 0; i < in_type->size(); i++) { auto eid = idx.entry_id(input_nids[i], 0); (*in_type)[i] = dtypes[eid]; } @@ -455,14 +494,13 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, imperative::CheckAndInferStorageType(g.get(), std::move(dev_masks), std::move(storage_type_inputs), true); - size_t num_input_arrays = 1; const auto& stypes = g->GetAttr("storage_type"); // Inferring the storage in the subgraph may infer the storage of the inputs. // We need to copy the inferred input storage back. const auto &input_nids = idx.input_nodes(); CHECK_EQ(input_nids.size(), in_attrs->size()); - for (size_t i = num_input_arrays; i < in_attrs->size(); i++) { + for (size_t i = 0; i < in_attrs->size(); i++) { auto eid = idx.entry_id(input_nids[i], 0); (*in_attrs)[i] = stypes[eid]; } @@ -526,8 +564,8 @@ NNVM_REGISTER_OP(_foreach) .set_attr("FStatefulComputeEx", ForeachComputeExCPU) .set_attr("key_var_num_args", "num_args") .add_argument("fn", "Symbol", "Input graph.") -.add_argument("input", "NDArray-or-Symbol", "The input array where we iterate over.") -.add_argument("states", "NDArray-or-Symbol[]", "The list of initial states.") +.add_argument("inputs", "NDArray-or-Symbol[]", + "The input arrays that include data arrays and states.") .add_arguments(ForeachParam::__FIELDS__()); NNVM_REGISTER_OP(_backward_foreach) From 9b4ac3c4da81670063f2e6329e6dd67cebd26cbd Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Thu, 10 May 2018 23:21:32 +0000 Subject: [PATCH 37/71] fix tests. --- tests/python/unittest/test_operator.py | 140 ++++++++++++++++--------- 1 file changed, 93 insertions(+), 47 deletions(-) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 6124fe093770..d83743217e32 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5663,53 +5663,105 @@ def test_float16_min_max(): assert np.finfo('float16').max == mx.nd.max(a).asscalar() +# TODO Test cases: +# in an iteration, data is stored in any location. +# # iterations: odd or even. +# multiple inputs and multiple outputs. +# test nested loop. + + @with_seed() def test_foreach(): - v3 = mx.sym.var("v3") - v4 = mx.sym.var("v4") - v5 = mx.sym.var("v5") + v3 = mx.sym.var("v0") + v4 = mx.sym.var("v1") + v5 = mx.sym.var("v2") # This tests foreach with accumulation sum. - def step(in1, states): - out = in1 + states[0] + v5 + def step1_sym(in1, states): + out = in1 * 2 + states[0] + v5 + return (out, [out]) + def step2_sym(in1, states): + out = states[0] + in1 * 2 + v5 return (out, [out]) - out = mx.sym.contrib.foreach(step, v3, [v4]) - out1 = out[0] * 2 - out = mx.sym.Group([out1, out[1][0]]) - arr1 = mx.nd.random.uniform(shape=(2, 2)) - arr2 = mx.nd.random.uniform(shape=(2)) - arr3 = mx.nd.random.uniform(shape=(2)) - arr_grad1 = mx.nd.empty(arr1.shape) - arr_grad2 = mx.nd.empty(arr2.shape) - arr_grad3 = mx.nd.empty(arr3.shape) - e = out.bind(ctx=mx.cpu(), args={'v3': arr1, 'v4': arr2, 'v5': arr3}, - args_grad={'v3': arr_grad1, 'v4': arr_grad2, 'v5': arr_grad3}) - e.forward(is_train=True) - # backward - out_grad = mx.nd.random.uniform(-10, 10, arr1.shape) - state_grad = mx.nd.random.uniform(-10, 10, arr2.shape) - e.backward([out_grad, state_grad]) - - res = [] - arr1.attach_grad() - arr2.attach_grad() - arr3.attach_grad() - with mx.autograd.record(): - for i in range(arr1.shape[0]): - if (i == 0): - tmp_res = mx.nd.expand_dims(arr2, 0) + mx.nd.expand_dims(arr1[i], 0) + mx.nd.expand_dims(arr3, 0) - else: - tmp_res = res[len(res) - 1] + mx.nd.expand_dims(arr1[i], 0) + mx.nd.expand_dims(arr3, 0) - res.append(tmp_res) - res1 = mx.nd.concat(*res, dim=0) - res2 = res1 * 2 - res = mx.nd.concat(res2, tmp_res, dim=0) - res.backward(mx.nd.concat(out_grad, mx.nd.expand_dims(state_grad, 0), dim=0)) - assert_almost_equal(e.outputs[0].asnumpy(), res2.asnumpy(), rtol=0.001, atol=0.0001) - assert_almost_equal(arr1.grad.asnumpy(), e.grad_arrays[0].asnumpy()) - assert_almost_equal(arr2.grad.asnumpy(), e.grad_arrays[1].asnumpy()) - assert_almost_equal(arr3.grad.asnumpy(), e.grad_arrays[2].asnumpy()) + def step1(data, state, free): + return data * 2 + state + free + + def step2(data, state, free): + return state + data * 2 + free + + def verify_foreach(step_sym, step, in_arrs, init_states, out_grads): + out = mx.sym.contrib.foreach(step_sym, v3, [v4]) + out1 = out[0] * 1 + out = mx.sym.Group([out1, out[1][0]]) + arr_grads = [] + arg_dict = {} + arg_grad_dict = {} + i = 0 + for arr in in_arrs: + arr_grad = mx.nd.empty(arr.shape) + arr_grads.append(arr_grad) + arg_dict['v'+str(i)] = arr + arg_grad_dict['v'+str(i)] = arr_grad + i = i + 1 + for arr in init_states: + arr_grad = mx.nd.empty(arr.shape) + arr_grads.append(arr_grad) + arg_dict['v'+str(i)] = arr + arg_grad_dict['v'+str(i)] = arr_grad + i = i + 1 + + gin_order = [] + for name in out.list_inputs(): + name = name[1:] + gin_order.append(int(name)) + + e = out.bind(ctx=mx.cpu(), args=arg_dict, args_grad=arg_grad_dict) + e.forward(is_train=True) + # backward + e.backward(out_grads) + + res = [] + for arr in in_arrs: + arr.attach_grad() + with mx.autograd.record(): + states = [mx.nd.expand_dims(s, 0) for s in init_states] + for i in range(in_arrs[0].shape[0]): + tmp_res = step(mx.nd.expand_dims(in_arrs[0][i], 0), + states[0], states[1]) + res.append(tmp_res) + states[0] = tmp_res + res1 = mx.nd.concat(*res, dim=0) + res2 = res1 * 1 + res = mx.nd.concat(res2, tmp_res, dim=0) + res.backward(mx.nd.concat(out_grads[0], mx.nd.expand_dims(out_grads[1], 0), dim=0)) + assert_almost_equal(e.outputs[0].asnumpy(), res2.asnumpy(), rtol=0.001, atol=0.0001) + for i in range(len(in_arrs)): + assert_almost_equal(in_arrs[i].grad.asnumpy(), e.grad_arrays[gin_order[i]].asnumpy()) + + # Test foreach with data in different locations among inputs, + # different numbers of iterations. + arrs = [mx.nd.random.uniform(shape=(2, 2))] + states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] + out_grads = [mx.nd.random.uniform(-10, 10, arrs[0].shape), + mx.nd.random.uniform(-10, 10, states[0].shape)] + verify_foreach(step1_sym, step1, arrs, states, out_grads) + + arrs = [mx.nd.random.uniform(shape=(3, 2))] + out_grads = [mx.nd.random.uniform(-10, 10, arrs[0].shape), + mx.nd.random.uniform(-10, 10, states[0].shape)] + verify_foreach(step1_sym, step1, arrs, states, out_grads) + + arrs = [mx.nd.random.uniform(shape=(2, 2))] + states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] + out_grads = [mx.nd.random.uniform(-10, 10, arrs[0].shape), + mx.nd.random.uniform(-10, 10, states[0].shape)] + verify_foreach(step2_sym, step2, arrs, states, out_grads) + + arrs = [mx.nd.random.uniform(shape=(3, 2))] + out_grads = [mx.nd.random.uniform(-10, 10, arrs[0].shape), + mx.nd.random.uniform(-10, 10, states[0].shape)] + verify_foreach(step2_sym, step2, arrs, states, out_grads) @with_seed() @@ -5807,12 +5859,6 @@ def sym_group(out): assert_almost_equal(e1.grad_arrays[i].asnumpy(), e2.grad_arrays[i].asnumpy()) -# TODO Test cases: -# in an iteration, data is stored in any location. -# # iterations: odd or even. -# multiple inputs and multiple outputs. - - @with_seed() def test_squeeze_op(): def check_squeeze_op(shape, axis=None): From a08de42761309466b9a33f61c8121bf2602daeea Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 11 May 2018 17:17:37 +0000 Subject: [PATCH 38/71] update tests. --- tests/python/unittest/test_operator.py | 149 ++++++++++++++++--------- 1 file changed, 98 insertions(+), 51 deletions(-) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index d83743217e32..4f1750ac3a2b 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -24,7 +24,7 @@ import itertools from numpy.testing import assert_allclose, assert_array_equal from mxnet.test_utils import * -from mxnet.base import py_str, MXNetError +from mxnet.base import py_str, MXNetError, _as_list from common import setup_module, with_seed import unittest @@ -5675,30 +5675,35 @@ def test_foreach(): v3 = mx.sym.var("v0") v4 = mx.sym.var("v1") v5 = mx.sym.var("v2") + v6 = mx.sym.var("v3") + v7 = mx.sym.var("v4") # This tests foreach with accumulation sum. - def step1_sym(in1, states): - out = in1 * 2 + states[0] + v5 + def step1(in1, states, free): + out = in1 * 2 + states[0] + free[0] return (out, [out]) - def step2_sym(in1, states): - out = states[0] + in1 * 2 + v5 + def step2(in1, states, free): + out = states[0] + in1 * 2 + free[0] return (out, [out]) - - def step1(data, state, free): - return data * 2 + state + free - - def step2(data, state, free): - return state + data * 2 + free - - def verify_foreach(step_sym, step, in_arrs, init_states, out_grads): - out = mx.sym.contrib.foreach(step_sym, v3, [v4]) - out1 = out[0] * 1 - out = mx.sym.Group([out1, out[1][0]]) + def step3(in1, states, free): + out = in1[0] + in1[1] + states[0] + states[1] + free[0] + return ([out, out * 2], [out * 2, out * 3]) + + def verify_foreach(step, in_syms, state_syms, free_syms, + in_arrs, init_states, frees, out_grads): + step_sym = lambda in_syms, state_syms : step(in_syms, state_syms, free_syms) + step_imp = lambda in_arrs, state_arrs : step(in_arrs, state_arrs, frees) + out = mx.sym.contrib.foreach(step_sym, in_syms, state_syms) + out1 = _as_list(out[0]) + for i in range(len(out1)): + out1[i] = out1[i] * 2 + out1.extend(out[1]) + out = mx.sym.Group(out1) arr_grads = [] arg_dict = {} arg_grad_dict = {} i = 0 - for arr in in_arrs: + for arr in _as_list(in_arrs): arr_grad = mx.nd.empty(arr.shape) arr_grads.append(arr_grad) arg_dict['v'+str(i)] = arr @@ -5710,6 +5715,12 @@ def verify_foreach(step_sym, step, in_arrs, init_states, out_grads): arg_dict['v'+str(i)] = arr arg_grad_dict['v'+str(i)] = arr_grad i = i + 1 + for arr in frees: + arr_grad = mx.nd.empty(arr.shape) + arr_grads.append(arr_grad) + arg_dict['v'+str(i)] = arr + arg_grad_dict['v'+str(i)] = arr_grad + i = i + 1 gin_order = [] for name in out.list_inputs(): @@ -5719,49 +5730,85 @@ def verify_foreach(step_sym, step, in_arrs, init_states, out_grads): e = out.bind(ctx=mx.cpu(), args=arg_dict, args_grad=arg_grad_dict) e.forward(is_train=True) # backward - e.backward(out_grads) + tmp_grads = out_grads[0][:] + tmp_grads.extend(out_grads[1]) + e.backward(tmp_grads) + # Below we use imperative to reimplement foreach and compute its gradients. res = [] - for arr in in_arrs: + for i in range(len(_as_list(out_grads[0]))): + res.append([]) + for arr in _as_list(in_arrs): + arr.attach_grad() + for arr in init_states: + arr.attach_grad() + for arr in frees: arr.attach_grad() with mx.autograd.record(): states = [mx.nd.expand_dims(s, 0) for s in init_states] - for i in range(in_arrs[0].shape[0]): - tmp_res = step(mx.nd.expand_dims(in_arrs[0][i], 0), - states[0], states[1]) - res.append(tmp_res) - states[0] = tmp_res - res1 = mx.nd.concat(*res, dim=0) - res2 = res1 * 1 - res = mx.nd.concat(res2, tmp_res, dim=0) - res.backward(mx.nd.concat(out_grads[0], mx.nd.expand_dims(out_grads[1], 0), dim=0)) - assert_almost_equal(e.outputs[0].asnumpy(), res2.asnumpy(), rtol=0.001, atol=0.0001) - for i in range(len(in_arrs)): - assert_almost_equal(in_arrs[i].grad.asnumpy(), e.grad_arrays[gin_order[i]].asnumpy()) + if isinstance(in_arrs, list): + num_iters = in_arrs[0].shape[0] + else: + num_iters = in_arrs.shape[0] + + for i in range(num_iters): + if isinstance(in_arrs, list): + data = [mx.nd.expand_dims(arr[i], 0) for arr in in_arrs] + else: + data = mx.nd.expand_dims(in_arrs[i], 0) + tmp_res = step_imp(data, states) + tmp_res1 = _as_list(tmp_res[0]) + for i in range(len(tmp_res1)): + res[i].append(tmp_res1[i]) + states = tmp_res[1] + res2 = [] + for l in res: + res2.append(mx.nd.concat(*l, dim=0) * 2) + tmp_res2 = res2[:] + tmp_res2.extend(tmp_res[1]) + res = mx.nd.concat(*tmp_res2, dim=0) + + tmp_grads = out_grads[0][:] + tmp_grads1 = [mx.nd.expand_dims(grad, 0) for grad in out_grads[1]] + tmp_grads.extend(tmp_grads1) + res.backward(mx.nd.concat(*tmp_grads, dim=0)) + for i in range(len(res2)): + assert_almost_equal(e.outputs[i].asnumpy(), res2[i].asnumpy(), rtol=0.001, atol=0.0001) + all_ins = _as_list(in_arrs) + all_ins.extend(init_states) + all_ins.extend(frees) + for i in range(len(all_ins)): + assert_almost_equal(all_ins[i].grad.asnumpy(), e.grad_arrays[gin_order[i]].asnumpy()) # Test foreach with data in different locations among inputs, # different numbers of iterations. - arrs = [mx.nd.random.uniform(shape=(2, 2))] - states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] - out_grads = [mx.nd.random.uniform(-10, 10, arrs[0].shape), - mx.nd.random.uniform(-10, 10, states[0].shape)] - verify_foreach(step1_sym, step1, arrs, states, out_grads) - - arrs = [mx.nd.random.uniform(shape=(3, 2))] - out_grads = [mx.nd.random.uniform(-10, 10, arrs[0].shape), - mx.nd.random.uniform(-10, 10, states[0].shape)] - verify_foreach(step1_sym, step1, arrs, states, out_grads) - - arrs = [mx.nd.random.uniform(shape=(2, 2))] + states = [mx.nd.random.uniform(shape=(2))] + frees = [mx.nd.random.uniform(shape=(2))] + arrs = mx.nd.random.uniform(shape=(2, 2)) + out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], + [mx.nd.random.uniform(-10, 10, states[0].shape)]] + verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads) + + arrs = mx.nd.random.uniform(shape=(3, 2)) + out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], + [mx.nd.random.uniform(-10, 10, states[0].shape)]] + verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads) + + arrs = mx.nd.random.uniform(shape=(2, 2)) + out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], + [mx.nd.random.uniform(-10, 10, states[0].shape)]] + verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads) + + arrs = mx.nd.random.uniform(shape=(3, 2)) + out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], + [mx.nd.random.uniform(-10, 10, states[0].shape)]] + verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads) + + arrs = [mx.nd.random.uniform(shape=(3, 2)), mx.nd.random.uniform(shape=(3, 2))] states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] - out_grads = [mx.nd.random.uniform(-10, 10, arrs[0].shape), - mx.nd.random.uniform(-10, 10, states[0].shape)] - verify_foreach(step2_sym, step2, arrs, states, out_grads) - - arrs = [mx.nd.random.uniform(shape=(3, 2))] - out_grads = [mx.nd.random.uniform(-10, 10, arrs[0].shape), - mx.nd.random.uniform(-10, 10, states[0].shape)] - verify_foreach(step2_sym, step2, arrs, states, out_grads) + out_grads = [[mx.nd.random.uniform(-10, 10, arrs[0].shape), mx.nd.random.uniform(-10, 10, arrs[1].shape)], + [mx.nd.random.uniform(-10, 10, states[0].shape), mx.nd.random.uniform(-10, 10, states[1].shape)]] + verify_foreach(step3, [v3, v4], [v5, v6], [v7], arrs, states, frees, out_grads) @with_seed() From 5c8a1876596ff40e01b7c80e45f9195b4d7441ce Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 12 May 2018 01:28:35 +0000 Subject: [PATCH 39/71] check state shape. --- src/operator/nn/control_flow.cc | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 6fe3675dfa5d..dfb00ac79c40 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -157,6 +157,8 @@ void ForeachState::Forward(std::vector cinputs, arg_names, params); // TODO here we only changed the output arrays in the arguments. // Will this be a problem? + // TODO we need to avoid shape inference and memory plan whenever the op is + // called. op->Forward(nullptr, inputs, outputs); if (is_recording) { @@ -442,6 +444,11 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, uint32_t eid = idx.entry_id(g->outputs[i]); (*out_shape)[i] = shapes[eid]; } + size_t num_states = g->outputs.size() - params.num_out_data; + for (size_t i = 0; i < num_states; i++) { + size_t loc = params.in_state_locs[i]; + CHECK((*out_shape)[i + params.num_out_data] == (*in_shape)[loc]); + } return true; } From e6b53bc4d73394488f64c5789b22ab0eb57f64a6 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 14 May 2018 23:39:04 +0000 Subject: [PATCH 40/71] enable nested foreach. --- src/imperative/imperative_utils.h | 18 ++++- tests/python/unittest/test_operator.py | 101 ++++++++++++++++++++----- 2 files changed, 94 insertions(+), 25 deletions(-) diff --git a/src/imperative/imperative_utils.h b/src/imperative/imperative_utils.h index 35c8bd4d2bf4..0051becc38f0 100644 --- a/src/imperative/imperative_utils.h +++ b/src/imperative/imperative_utils.h @@ -464,12 +464,18 @@ inline void PushOperator(const OpStatePtr& state, InvalidateOutputs(outputs, req); #endif fcompute_ex(state, opctx, inputs, req, outputs); - if (ctx.dev_mask() == gpu::kDevMask && exec_type == ExecType::kSync) { + if (ctx.dev_mask() == gpu::kDevMask && exec_type == ExecType::kSync + && rctx.get_stream()) { rctx.get_stream()->Wait(); } }; - if (exec_type == ExecType::kSync) { + // For operators with subgraphs, we need to invoke them in the main thread + // instead of the threaded engine. + if (!attrs.subgraphs.empty()) { + RunContext rctx{ctx, nullptr}; + run(rctx, engine::CallbackOnComplete()); + } else if (exec_type == ExecType::kSync) { Engine::Get()->PushSync( [=](RunContext rctx) { run(rctx, engine::CallbackOnComplete()); }, ctx, read_vars, write_vars, FnProperty::kNormal, 0, @@ -509,12 +515,16 @@ inline void PushOperator(const OpStatePtr& state, fcompute(state, opctx, input_blobs, tmp_req, output_blobs); // post-fcompute fallback, cast to original storage type, if necessary CastNonDefaultStorage(post_temp_src, post_temp_dst, opctx, is_gpu); - if (is_gpu && exec_type == ExecType::kSync) { + if (is_gpu && exec_type == ExecType::kSync + && rctx.get_stream()) { rctx.get_stream()->Wait(); } }; - if (exec_type == ExecType::kSync) { + if (!attrs.subgraphs.empty()) { + RunContext rctx{ctx, nullptr}; + run(rctx, engine::CallbackOnComplete()); + } else if (exec_type == ExecType::kSync) { Engine::Get()->PushSync( [=](RunContext rctx) { run(rctx, engine::CallbackOnComplete()); diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 4f1750ac3a2b..9e6a9710a5e9 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5663,13 +5663,6 @@ def test_float16_min_max(): assert np.finfo('float16').max == mx.nd.max(a).asscalar() -# TODO Test cases: -# in an iteration, data is stored in any location. -# # iterations: odd or even. -# multiple inputs and multiple outputs. -# test nested loop. - - @with_seed() def test_foreach(): v3 = mx.sym.var("v0") @@ -5690,7 +5683,7 @@ def step3(in1, states, free): return ([out, out * 2], [out * 2, out * 3]) def verify_foreach(step, in_syms, state_syms, free_syms, - in_arrs, init_states, frees, out_grads): + in_arrs, init_states, frees, out_grads, is_train=True): step_sym = lambda in_syms, state_syms : step(in_syms, state_syms, free_syms) step_imp = lambda in_arrs, state_arrs : step(in_arrs, state_arrs, frees) out = mx.sym.contrib.foreach(step_sym, in_syms, state_syms) @@ -5727,12 +5720,14 @@ def verify_foreach(step, in_syms, state_syms, free_syms, name = name[1:] gin_order.append(int(name)) - e = out.bind(ctx=mx.cpu(), args=arg_dict, args_grad=arg_grad_dict) - e.forward(is_train=True) - # backward - tmp_grads = out_grads[0][:] - tmp_grads.extend(out_grads[1]) - e.backward(tmp_grads) + e = out.bind(ctx=mx.cpu(), args=arg_dict, args_grad=arg_grad_dict, + ) + e.forward(is_train=is_train) + if (is_train): + # backward + tmp_grads = out_grads[0][:] + tmp_grads.extend(out_grads[1]) + e.backward(tmp_grads) # Below we use imperative to reimplement foreach and compute its gradients. res = [] @@ -5771,14 +5766,26 @@ def verify_foreach(step, in_syms, state_syms, free_syms, tmp_grads = out_grads[0][:] tmp_grads1 = [mx.nd.expand_dims(grad, 0) for grad in out_grads[1]] tmp_grads.extend(tmp_grads1) - res.backward(mx.nd.concat(*tmp_grads, dim=0)) + if (is_train): + res.backward(mx.nd.concat(*tmp_grads, dim=0)) for i in range(len(res2)): - assert_almost_equal(e.outputs[i].asnumpy(), res2[i].asnumpy(), rtol=0.001, atol=0.0001) - all_ins = _as_list(in_arrs) - all_ins.extend(init_states) - all_ins.extend(frees) - for i in range(len(all_ins)): - assert_almost_equal(all_ins[i].grad.asnumpy(), e.grad_arrays[gin_order[i]].asnumpy()) + assert_almost_equal(e.outputs[i].asnumpy(), res2[i].asnumpy(), + rtol=0.001, atol=0.0001) + if (is_train): + all_ins = _as_list(in_arrs) + all_ins.extend(init_states) + all_ins.extend(frees) + for i in range(len(all_ins)): + assert_almost_equal(all_ins[i].grad.asnumpy(), + e.grad_arrays[gin_order[i]].asnumpy()) + + # Test cases: + # * graph inputs are stored in different orders. + # This is to test if foreach finds the data arrays and weight arrays + # in the right location. + # * the number of iterations: odd or even. + # * multiple inputs and multiple outputs. + # * inference. # Test foreach with data in different locations among inputs, # different numbers of iterations. @@ -5788,27 +5795,79 @@ def verify_foreach(step, in_syms, state_syms, free_syms, out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], [mx.nd.random.uniform(-10, 10, states[0].shape)]] verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads) + verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads, False) arrs = mx.nd.random.uniform(shape=(3, 2)) out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], [mx.nd.random.uniform(-10, 10, states[0].shape)]] verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads) + verify_foreach(step1, v3, [v4], [v5], arrs, states, frees, out_grads, False) arrs = mx.nd.random.uniform(shape=(2, 2)) out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], [mx.nd.random.uniform(-10, 10, states[0].shape)]] verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads) + verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads, False) arrs = mx.nd.random.uniform(shape=(3, 2)) out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], [mx.nd.random.uniform(-10, 10, states[0].shape)]] verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads) + verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads, False) arrs = [mx.nd.random.uniform(shape=(3, 2)), mx.nd.random.uniform(shape=(3, 2))] states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] out_grads = [[mx.nd.random.uniform(-10, 10, arrs[0].shape), mx.nd.random.uniform(-10, 10, arrs[1].shape)], [mx.nd.random.uniform(-10, 10, states[0].shape), mx.nd.random.uniform(-10, 10, states[1].shape)]] verify_foreach(step3, [v3, v4], [v5, v6], [v7], arrs, states, frees, out_grads) + verify_foreach(step3, [v3, v4], [v5, v6], [v7], arrs, states, frees, out_grads, False) + + +@with_seed() +def test_foreach_nested(): + # Test nested foreach. + def step_in(in1, states): + out = in1 * 2 + states[0] + return (out, [out]) + + def step(in1, states): + out1 = mx.sym.contrib.foreach(step_in, in1, states) + out = mx.sym.broadcast_add(out1[0], states[0]) + return (out, [mx.sym.squeeze(mx.sym.slice(out, begin=(0, 0), end=(1, 2)))]) + + data_sym = mx.sym.var("v1") + state_sym = mx.sym.var("v2") + out = mx.sym.contrib.foreach(step, data_sym, [state_sym]) + + out1 = _as_list(out[0]) + for i in range(len(out1)): + out1[i] = out1[i] + out1.extend(out[1]) + out = mx.sym.Group(out1) + + data = mx.nd.arange(4).reshape((1, 2, 2)) + state = mx.nd.arange(2) + data_grad = mx.nd.empty(data.shape) + state_grad = mx.nd.empty(state.shape) + e = out.bind(ctx=mx.cpu(), args={'v1':data, 'v2':state}, + args_grad={'v1':data_grad, 'v2':state_grad}) + e.forward(is_train=True) + out = mx.nd.zeros_like(data) + for i in range(data.shape[0]): + data1 = data[i] + out1 = mx.nd.zeros_like(data1) + for j in range(data1.shape[0]): + if (j > 0): + out1[j] = out1[j-1] + data1[j] * 2 + else: + out1[j] = data1[j] * 2 + state + if (i > 0): + state = mx.nd.squeeze(mx.nd.slice(out[i-1], begin=(0, 0), end=(1, 2))) + out[i] = mx.nd.broadcast_add(out1, state) + else: + out[i] = mx.nd.broadcast_add(out1, state) + out = out + assert_almost_equal(out.asnumpy(), e.outputs[0].asnumpy(), rtol=0.001, atol=0.0001) @with_seed() From 65c8515ff55952ff43fb3e9897ed3c6f273dcee2 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 16 May 2018 23:09:33 +0000 Subject: [PATCH 41/71] remove print. --- src/c_api/c_api_symbolic.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index a2ebefc8255d..b0d2ddf3b4b9 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -362,7 +362,6 @@ int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **out_arr, int *out_s nnvm::Symbol *s = new nnvm::Symbol(); s->outputs.push_back(e); input_syms.push_back(s); - std::cout << p->attrs.name << std::endl; } } } From 7e05c980a7b45f0a35d083093d964adfef138e1b Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Thu, 17 May 2018 14:32:25 +0000 Subject: [PATCH 42/71] fix a bug in test. --- tests/python/unittest/test_operator.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 9e6a9710a5e9..3ed44c80df04 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5772,7 +5772,7 @@ def verify_foreach(step, in_syms, state_syms, free_syms, assert_almost_equal(e.outputs[i].asnumpy(), res2[i].asnumpy(), rtol=0.001, atol=0.0001) if (is_train): - all_ins = _as_list(in_arrs) + all_ins = _as_list(in_arrs)[:] all_ins.extend(init_states) all_ins.extend(frees) for i in range(len(all_ins)): @@ -5787,8 +5787,6 @@ def verify_foreach(step, in_syms, state_syms, free_syms, # * multiple inputs and multiple outputs. # * inference. - # Test foreach with data in different locations among inputs, - # different numbers of iterations. states = [mx.nd.random.uniform(shape=(2))] frees = [mx.nd.random.uniform(shape=(2))] arrs = mx.nd.random.uniform(shape=(2, 2)) @@ -5815,6 +5813,7 @@ def verify_foreach(step, in_syms, state_syms, free_syms, verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads) verify_foreach(step2, v3, [v4], [v5], arrs, states, frees, out_grads, False) + # Test multiple inputs and outputs. arrs = [mx.nd.random.uniform(shape=(3, 2)), mx.nd.random.uniform(shape=(3, 2))] states = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] out_grads = [[mx.nd.random.uniform(-10, 10, arrs[0].shape), mx.nd.random.uniform(-10, 10, arrs[1].shape)], From fa2b9bde6cf6a43af87340416a3045258de02d72 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 18 May 2018 06:18:28 +0000 Subject: [PATCH 43/71] handle infer storage type for backward. --- src/operator/nn/control_flow.cc | 96 +++++++++++++++++++++++++++------ 1 file changed, 81 insertions(+), 15 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index dfb00ac79c40..87fb48cfb749 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -157,8 +157,10 @@ void ForeachState::Forward(std::vector cinputs, arg_names, params); // TODO here we only changed the output arrays in the arguments. // Will this be a problem? - // TODO we need to avoid shape inference and memory plan whenever the op is - // called. + // TODO(zhengda) we need to avoid shape inference and memory plan whenever the op is + // called. Currently, CachedOp allocates memory each time Forward is called. + // I need to fix this once the PR for static memory allocation in CachedOp is + // merged. https://github.com/apache/incubator-mxnet/pull/10817 op->Forward(nullptr, inputs, outputs); if (is_recording) { @@ -410,7 +412,6 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, CHECK_EQ(attrs.subgraphs.size(), 1U); auto g = std::make_shared(); g->outputs = attrs.subgraphs[0]->outputs; - // TODO(zhengda) We should avoid creating an index graph so many times. const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_shape->size()); CHECK_EQ(idx.outputs().size(), out_shape->size()); @@ -460,7 +461,6 @@ static bool ForeachType(const nnvm::NodeAttrs& attrs, CHECK_EQ(attrs.subgraphs.size(), 1U); auto g = std::make_shared(); g->outputs = attrs.subgraphs[0]->outputs; - // TODO(zhengda) We should avoid creating an index graph so many times. const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_type->size()); CHECK_EQ(idx.outputs().size(), out_type->size()); @@ -492,7 +492,6 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, CHECK_EQ(attrs.subgraphs.size(), 1U); auto g = std::make_shared(); g->outputs = attrs.subgraphs[0]->outputs; - // TODO(zhengda) We should avoid creating an index graph so many times. const auto& idx = g->indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_attrs->size()); CHECK_EQ(idx.outputs().size(), out_attrs->size()); @@ -525,9 +524,69 @@ static bool BackwardForeachStorageType(const nnvm::NodeAttrs& attrs, DispatchMode* dispatch_mode, std::vector *in_attrs, std::vector *out_attrs) { - // TODO I need to set storage type properly. - return storage_type_assign(out_attrs, mxnet::kDefaultStorage, - dispatch_mode, DispatchMode::kFComputeEx); + using namespace nnvm; + const ForeachParam& params = nnvm::get(attrs.parsed); + CHECK_EQ(out_attrs->size(), (size_t) params.num_args - 1); + CHECK_EQ(attrs.subgraphs.size(), 1U); + // construct backward graph + nnvm::Graph grad_graph; + nnvm::Graph fwd_graph; + std::vector potential_nodes; + { + fwd_graph.outputs = attrs.subgraphs[0]->outputs; + std::vector ograd_entries; + ograd_entries.reserve(fwd_graph.outputs.size()); + for (size_t i = 0; i < fwd_graph.outputs.size(); ++i) { + ograd_entries.emplace_back(NodeEntry{Node::Create(), 0, 0}); + } + nnvm::Symbol subgraph_sym = *attrs.subgraphs[0]; + + std::vector xs; + std::vector args = subgraph_sym.ListInputs(nnvm::Symbol::kReadOnlyArgs); + xs.reserve(args.size()); + for (const auto& i : args) + xs.emplace_back(NodeEntry{i, 0, 0}); + CHECK_GT(xs.size(), 0) + << "There are no inputs in computation graph that require gradients."; + + static const std::vector zero_ops{Op::Get("zeros_like"), Op::Get("_zeros")}; + grad_graph = pass::Gradient( + fwd_graph, fwd_graph.outputs, xs, ograd_entries, + exec::AggregateGradient, nullptr, nullptr, + zero_ops, "_copy"); + potential_nodes.reserve(fwd_graph.outputs.size() + xs.size() + ograd_entries.size()); + for (auto e : ograd_entries) + potential_nodes.push_back(e.node.get()); + for (auto e : xs) + potential_nodes.push_back(e.node.get()); + for (auto e : fwd_graph.outputs) + potential_nodes.push_back(e.node.get()); + } + + const auto& idx = grad_graph.indexed_graph(); + auto input_nodes = idx.input_nodes(); + StorageTypeVector storage_type_inputs(input_nodes.size()); + for (size_t i = 0; i < input_nodes.size(); i++) { + auto node_id = input_nodes[i]; + const nnvm::IndexedGraph::Node &n = idx[node_id]; + auto it = std::find(potential_nodes.begin(), potential_nodes.end(), n.source); + CHECK(it != potential_nodes.end()); + size_t idx = it - potential_nodes.begin(); + CHECK_LT(idx, in_attrs->size()); + storage_type_inputs[i] = in_attrs->at(idx); + } + CHECK_EQ(idx.outputs().size(), out_attrs->size()); + exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); + imperative::CheckAndInferStorageType(&grad_graph, std::move(dev_masks), + std::move(storage_type_inputs), true); + + const auto& stypes = grad_graph.GetAttr("storage_type"); + *dispatch_mode = DispatchMode::kFComputeEx; + auto &outputs = idx.outputs(); + CHECK(outputs.size() == out_attrs->size()); + for (size_t i = 0; i < out_attrs->size(); i++) + (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; + return true; } static OpStatePtr CreateForeachState(const NodeAttrs& attrs, @@ -538,10 +597,12 @@ static OpStatePtr CreateForeachState(const NodeAttrs& attrs, return OpStatePtr::Create(*attrs.subgraphs[0], params); } -void ForeachParamParser(nnvm::NodeAttrs* attrs) { - ParamParser(attrs); - // This is to indicate that the operator has a subgraph. - attrs->subgraphs.resize(1); +static std::vector +ForeachGradient(const nnvm::NodePtr& n, const std::vector& ograds) { + ElemwiseGradUseInOut fgrad{"_backward_foreach"}; + std::vector entries = fgrad(n, ograds); + entries[0].node->attrs.subgraphs = n->attrs.subgraphs; + return entries; } NNVM_REGISTER_OP(_foreach) @@ -558,13 +619,18 @@ NNVM_REGISTER_OP(_foreach) }) .set_attr("FListInputNames", [](const NodeAttrs& attrs) { - return std::vector{"fn", "data1", "data2"}; + const ForeachParam& params = nnvm::get(attrs.parsed); + std::vector names; + names.push_back("fn"); + for (int i = 0; i < params.num_args - 1; i++) + names.push_back("data" + std::to_string(i)); + return names; }) .set_attr("FInputGraph", [](const NodeAttrs& attrs) { return std::vector{0}; }) -.set_attr("FGradient", ElemwiseGradUseInOut{"_backward_foreach"}) +.set_attr("FGradient", ForeachGradient) .set_attr("FCreateOpState", CreateForeachState) .set_attr("FInferShape", ForeachShape) .set_attr("FInferType", ForeachType) @@ -585,7 +651,7 @@ NNVM_REGISTER_OP(_backward_foreach) return params.num_args - 1; }) .set_attr("FInferStorageType", BackwardForeachStorageType) -.set_attr_parser(ForeachParamParser) +.set_attr_parser(ParamParser) .set_attr("TIsLayerOpBackward", true) .set_attr("TIsBackward", true) .set_attr("FStatefulComputeEx", ForeachGradComputeExCPU); From fcf4c3465eb03e176f5fc722d33ea75fa4a571b5 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 18 May 2018 06:31:48 +0000 Subject: [PATCH 44/71] address comments. --- python/mxnet/symbol/contrib.py | 31 ++++++++++++++++++------------- src/c_api/c_api_symbolic.cc | 2 +- src/executor/exec_pass.h | 4 +--- src/operator/nn/control_flow.cc | 4 ++-- 4 files changed, 22 insertions(+), 19 deletions(-) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index 1d82cdc92763..c236dffa63cd 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -111,14 +111,19 @@ def _get_graph_inputs(subg, name, prefix): syms.append(s) return syms -def foreach(func, input, init_states, back_prop=False, name="foreach"): +def foreach(func, data, init_states, back_prop=False, name="foreach"): assert isinstance(init_states, list), "init_states should be a list" states = [] + + # TODO(zhengda) If the input python function references to the symbols outside + # the python function, we need to prune the computation graph constructed from + # the function. One way of doing it is to mark the nodes in the computation graph + # with AttrScope and prune the nodes without the special attribute. with AttrScope(subgraph_name=name): - if isinstance(input, list): - in_eles = [symbol.var(sym.name) for sym in input] + if isinstance(data, list): + in_eles = [symbol.var(sym.name) for sym in data] else: - in_eles = symbol.var(input.name) + in_eles = symbol.var(data.name) for s in init_states: states.append(symbol.var(s.name)) @@ -132,21 +137,21 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): "the number of output states (%d) should be the same as input states (%d)" \ % (len(sym_out[1]), len(init_states)) - if (isinstance(sym_out[0], list)): + if isinstance(sym_out[0], list): flat_out = sym_out[0] else: flat_out = [sym_out[0]] num_out_data = len(flat_out) for s in sym_out[1]: # There is a problem if the outputs are the same as the inputs - # or the first output. - # TODO this is a temp fix. + # or the first output. By calling identity, we can make sure that + # all symbols will refer to different NDArrays. flat_out.append(symbol.op.identity(s)) g = symbol.Group(flat_out) input_syms = _get_graph_inputs(g, name, "ro_var") - if (isinstance(input, list)): - num_inputs = len(input) + if isinstance(data, list): + num_inputs = len(data) else: num_inputs = 1 @@ -161,7 +166,7 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): ordered_ins = [] states_map = {sym.name:sym for sym in init_states} state_names = states_map.keys() - data_syms = _as_list(input) + data_syms = _as_list(data) data_map = {sym.name:sym for sym in data_syms} data_names = data_map.keys() in_state_locs = [] @@ -169,10 +174,10 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): for in_name in g.list_inputs(): assert in_name in gin_names, "The input variable %s can't be found in graph inputs: %s" \ % (in_name, str(gin_names)) - if (in_name in state_names): + if in_name in state_names: ordered_ins.append(states_map[in_name]) in_state_locs.append(len(ordered_ins) - 1) - elif (in_name in data_names): + elif in_name in data_names: ordered_ins.append(data_map[in_name]) in_data_locs.append(len(ordered_ins) - 1) else: @@ -183,7 +188,7 @@ def foreach(func, input, init_states, back_prop=False, name="foreach"): ret = symbol._internal._foreach(g, *ordered_ins, num_outputs=num_outputs, num_out_data=num_out_data, in_state_locs=in_state_locs, in_data_locs=in_data_locs) - if (num_outputs - num_states > 1): + if num_outputs - num_states > 1: outs = [] for i in range(num_outputs - num_states): outs.append(ret[i]) diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index b0d2ddf3b4b9..3ffa07256266 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -368,7 +368,7 @@ int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **out_arr, int *out_s CHECK(input_syms.size() <= max_out_size); *out_size = input_syms.size(); memcpy(out_arr, input_syms.data(), sizeof(*out_arr) * input_syms.size()); - API_END(); + API_END_HANDLE_ERROR(); } int MXSymbolCreateFromFile(const char *fname, SymbolHandle *out) { diff --git a/src/executor/exec_pass.h b/src/executor/exec_pass.h index aa28e37500d5..f49fcf61db21 100644 --- a/src/executor/exec_pass.h +++ b/src/executor/exec_pass.h @@ -64,9 +64,7 @@ class OpExecutor { OpContext op_ctx; /*! \brief virtual destructor */ virtual ~OpExecutor() {} - virtual bool HasSubgraph() const { - return false; - } + virtual bool HasSubgraph() const = 0; /*! * \brief Setup the executor for given NDArray member * this can be called multiple times if NDArray changed during reshape. diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 87fb48cfb749..a454fb32d296 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -328,7 +328,7 @@ static void ForeachGradComputeExCPU(const OpStatePtr& state_ptr, std::unordered_set in_state_locs(params.in_state_locs.begin(), params.in_state_locs.end()); // The inputs contain out gradients, inputs and outputs. - size_t len = inputs[0].shape()[iter_dim]; + int len = inputs[0].shape()[iter_dim]; size_t num_output_data = params.num_out_data; // In backward computation, we need to run iterations from backwards. @@ -637,7 +637,7 @@ NNVM_REGISTER_OP(_foreach) .set_attr("FStatefulComputeEx", ForeachComputeExCPU) .set_attr("key_var_num_args", "num_args") .add_argument("fn", "Symbol", "Input graph.") -.add_argument("inputs", "NDArray-or-Symbol[]", +.add_argument("data", "NDArray-or-Symbol[]", "The input arrays that include data arrays and states.") .add_arguments(ForeachParam::__FIELDS__()); From 54f6efa802e9ff89c4182ca8ae99ef148d815b07 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 18 May 2018 06:34:26 +0000 Subject: [PATCH 45/71] address comments. --- src/operator/nn/control_flow.cc | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index a454fb32d296..60e7d5be55b6 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -56,26 +56,6 @@ struct ForeachParam : public dmlc::Parameter { DMLC_REGISTER_PARAMETER(ForeachParam); -// The input arguments are ordered in the following order: -// in, state0, state1, ... -// We need to reorder them in the same order as the input nodes of the subgraph. -template -static std::vector ReorderInputs(const std::vector &in, const nnvm::IndexedGraph& idx) { - std::vector ret(in.size()); - CHECK_EQ(idx.input_nodes().size(), in.size()); - for (size_t i = 0; i < idx.input_nodes().size(); i++) { - std::string name = idx[idx.input_nodes()[i]].source->attrs.name; - if (name == "in") { - ret[i] = in[0]; - } else { - auto idx_str = name.substr(5); - int idx = std::stoi(idx_str); - ret[i] = in[idx + 1]; - } - } - return ret; -} - class ForeachState { // These are output arrays from all iterations. // They also contain the Op state for each CachedOp. From a8f86f89eb9514c21697f88a421822d31a2a91a1 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 18 May 2018 16:51:21 +0000 Subject: [PATCH 46/71] move some common functions out. --- src/operator/nn/control_flow.cc | 115 +------------------ src/operator/nn/subgraph_op_common.cc | 154 ++++++++++++++++++++++++++ src/operator/nn/subgraph_op_common.h | 46 ++++++++ 3 files changed, 206 insertions(+), 109 deletions(-) create mode 100644 src/operator/nn/subgraph_op_common.cc create mode 100644 src/operator/nn/subgraph_op_common.h diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 60e7d5be55b6..ab3e4118cdb0 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -27,6 +27,7 @@ #include "../operator_common.h" #include "../elemwise_op_common.h" #include "../../imperative/imperative_utils.h" +#include "./subgraph_op_common.h" namespace mxnet { namespace op { @@ -437,29 +438,8 @@ static bool ForeachType(const nnvm::NodeAttrs& attrs, std::vector *in_type, std::vector *out_type) { const ForeachParam& params = nnvm::get(attrs.parsed); CHECK_EQ(out_type->size(), (size_t) params.num_outputs); - nnvm::DTypeVector dtype_inputs = *in_type; CHECK_EQ(attrs.subgraphs.size(), 1U); - auto g = std::make_shared(); - g->outputs = attrs.subgraphs[0]->outputs; - const auto& idx = g->indexed_graph(); - CHECK_EQ(idx.input_nodes().size(), in_type->size()); - CHECK_EQ(idx.outputs().size(), out_type->size()); - imperative::CheckAndInferType(g.get(), std::move(dtype_inputs), true); - - const auto &dtypes = g->GetAttr("dtype"); - - // Inferring the data type in the subgraph may infer the data type of the inputs. - // We need to copy the inferred input data types back. - const auto &input_nids = idx.input_nodes(); - CHECK_EQ(input_nids.size(), in_type->size()); - for (size_t i = 0; i < in_type->size(); i++) { - auto eid = idx.entry_id(input_nids[i], 0); - (*in_type)[i] = dtypes[eid]; - } - - for (size_t i = 0; i < g->outputs.size(); i++) - (*out_type)[i] = dtypes[idx.entry_id(g->outputs[i])]; - return true; + return InferSubgraphDataType(*attrs.subgraphs[0], in_type, out_type); } static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, @@ -470,33 +450,8 @@ static bool ForeachStorageType(const nnvm::NodeAttrs& attrs, const ForeachParam& params = nnvm::get(attrs.parsed); CHECK_EQ(out_attrs->size(), (size_t) params.num_outputs); CHECK_EQ(attrs.subgraphs.size(), 1U); - auto g = std::make_shared(); - g->outputs = attrs.subgraphs[0]->outputs; - const auto& idx = g->indexed_graph(); - CHECK_EQ(idx.input_nodes().size(), in_attrs->size()); - CHECK_EQ(idx.outputs().size(), out_attrs->size()); - exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); - StorageTypeVector storage_type_inputs = *in_attrs; - imperative::CheckAndInferStorageType(g.get(), std::move(dev_masks), - std::move(storage_type_inputs), true); - - const auto& stypes = g->GetAttr("storage_type"); - - // Inferring the storage in the subgraph may infer the storage of the inputs. - // We need to copy the inferred input storage back. - const auto &input_nids = idx.input_nodes(); - CHECK_EQ(input_nids.size(), in_attrs->size()); - for (size_t i = 0; i < in_attrs->size(); i++) { - auto eid = idx.entry_id(input_nids[i], 0); - (*in_attrs)[i] = stypes[eid]; - } - - *dispatch_mode = DispatchMode::kFComputeEx; - auto &outputs = idx.outputs(); - CHECK(outputs.size() == out_attrs->size()); - for (size_t i = 0; i < out_attrs->size(); i++) - (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; - return true; + return InferSubgraphStorage(*attrs.subgraphs[0], dev_mask, + dispatch_mode, in_attrs, out_attrs); } static bool BackwardForeachStorageType(const nnvm::NodeAttrs& attrs, @@ -504,69 +459,11 @@ static bool BackwardForeachStorageType(const nnvm::NodeAttrs& attrs, DispatchMode* dispatch_mode, std::vector *in_attrs, std::vector *out_attrs) { - using namespace nnvm; const ForeachParam& params = nnvm::get(attrs.parsed); CHECK_EQ(out_attrs->size(), (size_t) params.num_args - 1); CHECK_EQ(attrs.subgraphs.size(), 1U); - // construct backward graph - nnvm::Graph grad_graph; - nnvm::Graph fwd_graph; - std::vector potential_nodes; - { - fwd_graph.outputs = attrs.subgraphs[0]->outputs; - std::vector ograd_entries; - ograd_entries.reserve(fwd_graph.outputs.size()); - for (size_t i = 0; i < fwd_graph.outputs.size(); ++i) { - ograd_entries.emplace_back(NodeEntry{Node::Create(), 0, 0}); - } - nnvm::Symbol subgraph_sym = *attrs.subgraphs[0]; - - std::vector xs; - std::vector args = subgraph_sym.ListInputs(nnvm::Symbol::kReadOnlyArgs); - xs.reserve(args.size()); - for (const auto& i : args) - xs.emplace_back(NodeEntry{i, 0, 0}); - CHECK_GT(xs.size(), 0) - << "There are no inputs in computation graph that require gradients."; - - static const std::vector zero_ops{Op::Get("zeros_like"), Op::Get("_zeros")}; - grad_graph = pass::Gradient( - fwd_graph, fwd_graph.outputs, xs, ograd_entries, - exec::AggregateGradient, nullptr, nullptr, - zero_ops, "_copy"); - potential_nodes.reserve(fwd_graph.outputs.size() + xs.size() + ograd_entries.size()); - for (auto e : ograd_entries) - potential_nodes.push_back(e.node.get()); - for (auto e : xs) - potential_nodes.push_back(e.node.get()); - for (auto e : fwd_graph.outputs) - potential_nodes.push_back(e.node.get()); - } - - const auto& idx = grad_graph.indexed_graph(); - auto input_nodes = idx.input_nodes(); - StorageTypeVector storage_type_inputs(input_nodes.size()); - for (size_t i = 0; i < input_nodes.size(); i++) { - auto node_id = input_nodes[i]; - const nnvm::IndexedGraph::Node &n = idx[node_id]; - auto it = std::find(potential_nodes.begin(), potential_nodes.end(), n.source); - CHECK(it != potential_nodes.end()); - size_t idx = it - potential_nodes.begin(); - CHECK_LT(idx, in_attrs->size()); - storage_type_inputs[i] = in_attrs->at(idx); - } - CHECK_EQ(idx.outputs().size(), out_attrs->size()); - exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); - imperative::CheckAndInferStorageType(&grad_graph, std::move(dev_masks), - std::move(storage_type_inputs), true); - - const auto& stypes = grad_graph.GetAttr("storage_type"); - *dispatch_mode = DispatchMode::kFComputeEx; - auto &outputs = idx.outputs(); - CHECK(outputs.size() == out_attrs->size()); - for (size_t i = 0; i < out_attrs->size(); i++) - (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; - return true; + return InferSubgraphBackwardStorage(*attrs.subgraphs[0], dev_mask, + dispatch_mode, in_attrs, out_attrs); } static OpStatePtr CreateForeachState(const NodeAttrs& attrs, diff --git a/src/operator/nn/subgraph_op_common.cc b/src/operator/nn/subgraph_op_common.cc new file mode 100644 index 000000000000..7efd48e0eed0 --- /dev/null +++ b/src/operator/nn/subgraph_op_common.cc @@ -0,0 +1,154 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "./subgraph_op_common.h" +#include "../operator_common.h" +#include "../../imperative/imperative_utils.h" + +namespace mxnet { +namespace op { + +bool InferSubgraphDataType(nnvm::Symbol &subgraph, std::vector *in_type, + std::vector *out_type) { + nnvm::DTypeVector dtype_inputs = *in_type; + nnvm::Graph g; + g.outputs = subgraph.outputs; + const auto& idx = g.indexed_graph(); + CHECK_EQ(idx.input_nodes().size(), in_type->size()); + CHECK_EQ(idx.outputs().size(), out_type->size()); + imperative::CheckAndInferType(&g, std::move(dtype_inputs), true); + + const auto &dtypes = g.GetAttr("dtype"); + + // Inferring the data type in the subgraph may infer the data type of the inputs. + // We need to copy the inferred input data types back. + const auto &input_nids = idx.input_nodes(); + CHECK_EQ(input_nids.size(), in_type->size()); + for (size_t i = 0; i < in_type->size(); i++) { + auto eid = idx.entry_id(input_nids[i], 0); + (*in_type)[i] = dtypes[eid]; + } + + for (size_t i = 0; i < g.outputs.size(); i++) + (*out_type)[i] = dtypes[idx.entry_id(g.outputs[i])]; + return true; +} + +bool InferSubgraphStorage(nnvm::Symbol &subgraph, + const int dev_mask, + DispatchMode* dispatch_mode, + std::vector *in_attrs, + std::vector *out_attrs) { + nnvm::Graph g; + g.outputs = subgraph.outputs; + const auto& idx = g.indexed_graph(); + CHECK_EQ(idx.input_nodes().size(), in_attrs->size()); + CHECK_EQ(idx.outputs().size(), out_attrs->size()); + exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); + StorageTypeVector storage_type_inputs = *in_attrs; + imperative::CheckAndInferStorageType(&g, std::move(dev_masks), + std::move(storage_type_inputs), true); + + const auto& stypes = g.GetAttr("storage_type"); + + // Inferring the storage in the subgraph may infer the storage of the inputs. + // We need to copy the inferred input storage back. + const auto &input_nids = idx.input_nodes(); + CHECK_EQ(input_nids.size(), in_attrs->size()); + for (size_t i = 0; i < in_attrs->size(); i++) { + auto eid = idx.entry_id(input_nids[i], 0); + (*in_attrs)[i] = stypes[eid]; + } + + *dispatch_mode = DispatchMode::kFComputeEx; + auto &outputs = idx.outputs(); + CHECK(outputs.size() == out_attrs->size()); + for (size_t i = 0; i < out_attrs->size(); i++) + (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; + return true; +} + +bool InferSubgraphBackwardStorage(nnvm::Symbol &subgraph, + const int dev_mask, + DispatchMode* dispatch_mode, + std::vector *in_attrs, + std::vector *out_attrs) { + using namespace nnvm; + // construct backward graph + nnvm::Graph grad_graph; + nnvm::Graph fwd_graph; + std::vector potential_nodes; + { + fwd_graph.outputs = subgraph.outputs; + std::vector ograd_entries; + ograd_entries.reserve(fwd_graph.outputs.size()); + for (size_t i = 0; i < fwd_graph.outputs.size(); ++i) { + ograd_entries.emplace_back(NodeEntry{Node::Create(), 0, 0}); + } + + std::vector xs; + std::vector args = subgraph.ListInputs(nnvm::Symbol::kReadOnlyArgs); + xs.reserve(args.size()); + for (const auto& i : args) + xs.emplace_back(NodeEntry{i, 0, 0}); + CHECK_GT(xs.size(), 0) + << "There are no inputs in computation graph that require gradients."; + + static const std::vector zero_ops{Op::Get("zeros_like"), Op::Get("_zeros")}; + grad_graph = pass::Gradient( + fwd_graph, fwd_graph.outputs, xs, ograd_entries, + exec::AggregateGradient, nullptr, nullptr, + zero_ops, "_copy"); + potential_nodes.reserve(fwd_graph.outputs.size() + xs.size() + ograd_entries.size()); + for (auto e : ograd_entries) + potential_nodes.push_back(e.node.get()); + for (auto e : xs) + potential_nodes.push_back(e.node.get()); + for (auto e : fwd_graph.outputs) + potential_nodes.push_back(e.node.get()); + } + + const auto& idx = grad_graph.indexed_graph(); + auto input_nodes = idx.input_nodes(); + StorageTypeVector storage_type_inputs(input_nodes.size()); + for (size_t i = 0; i < input_nodes.size(); i++) { + auto node_id = input_nodes[i]; + const nnvm::IndexedGraph::Node &n = idx[node_id]; + auto it = std::find(potential_nodes.begin(), potential_nodes.end(), n.source); + CHECK(it != potential_nodes.end()); + size_t idx = it - potential_nodes.begin(); + CHECK_LT(idx, in_attrs->size()); + storage_type_inputs[i] = in_attrs->at(idx); + } + CHECK_EQ(idx.outputs().size(), out_attrs->size()); + exec::DevMaskVector dev_masks(idx.num_nodes(), dev_mask); + imperative::CheckAndInferStorageType(&grad_graph, std::move(dev_masks), + std::move(storage_type_inputs), true); + + const auto& stypes = grad_graph.GetAttr("storage_type"); + *dispatch_mode = DispatchMode::kFComputeEx; + auto &outputs = idx.outputs(); + CHECK(outputs.size() == out_attrs->size()); + for (size_t i = 0; i < out_attrs->size(); i++) + (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; + return true; +} + +} // namespace op +} // namespace mxnet diff --git a/src/operator/nn/subgraph_op_common.h b/src/operator/nn/subgraph_op_common.h new file mode 100644 index 000000000000..6412bb45e8a3 --- /dev/null +++ b/src/operator/nn/subgraph_op_common.h @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#ifndef MXNET_OPERATOR_NN_SUBGRAPH_OP_COMMON_H_ +#define MXNET_OPERATOR_NN_SUBGRAPH_OP_COMMON_H_ + +#include +#include +#include + +namespace mxnet { +namespace op { + +bool InferSubgraphDataType(nnvm::Symbol &subgraph, std::vector *in_type, + std::vector *out_type); +bool InferSubgraphStorage(nnvm::Symbol &subgraph, + const int dev_mask, + DispatchMode* dispatch_mode, + std::vector *in_attrs, + std::vector *out_attrs); +bool InferSubgraphBackwardStorage(nnvm::Symbol &subgraph, + const int dev_mask, + DispatchMode* dispatch_mode, + std::vector *in_attrs, + std::vector *out_attrs); + +} // namespace op +} // namespace mxnet + +#endif // MXNET_OPERATOR_NN_SUBGRAPH_OP_COMMON_H_ From 0f894faffa12ffb692b4d12d0ea1f6ff1f1a033d Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 18 May 2018 17:47:02 +0000 Subject: [PATCH 47/71] address comments. --- src/operator/nn/control_flow.cc | 26 +++++++++++++------------- src/operator/nn/subgraph_op_common.cc | 14 +++++++------- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index ab3e4118cdb0..8e242e950aff 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -391,14 +391,14 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, shape_inputs[loc] = TShape(in_shape->at(loc).begin() + 1, in_shape->at(loc).end()); } CHECK_EQ(attrs.subgraphs.size(), 1U); - auto g = std::make_shared(); - g->outputs = attrs.subgraphs[0]->outputs; - const auto& idx = g->indexed_graph(); + nnvm::Graph g; + g.outputs = attrs.subgraphs[0]->outputs; + const auto& idx = g.indexed_graph(); CHECK_EQ(idx.input_nodes().size(), in_shape->size()); CHECK_EQ(idx.outputs().size(), out_shape->size()); - imperative::CheckAndInferShape(g.get(), std::move(shape_inputs), true); - const auto& shapes = g->GetAttr("shape"); + imperative::CheckAndInferShape(&g, std::move(shape_inputs), true); + const auto& shapes = g.GetAttr("shape"); // Inferring the shape in the subgraph may infer the shape of the inputs. // We need to copy the inferred input shapes back. const auto &input_nids = idx.input_nodes(); @@ -407,26 +407,26 @@ static bool ForeachShape(const nnvm::NodeAttrs& attrs, auto eid = idx.entry_id(input_nids[i], 0); // If the input shape is none, we should update them. if ((*in_shape)[i].ndim() == 0 || (*in_shape)[i].Size() == 0) - (*in_shape)[i] = shapes[eid]; + SHAPE_ASSIGN_CHECK(*in_shape, i, shapes[eid]); } // For the shape of output data. for (int i = 0; i < params.num_out_data; i++) { - uint32_t eid = idx.entry_id(g->outputs[i]); + uint32_t eid = idx.entry_id(g.outputs[i]); const auto& g_out_shape = shapes[eid]; - auto &out = (*out_shape)[i]; - out = TShape(g_out_shape.ndim() + 1); + auto out = TShape(g_out_shape.ndim() + 1); out[0] = len; for (size_t i = 1; i < out.ndim(); i++) out[i] = g_out_shape[i - 1]; + SHAPE_ASSIGN_CHECK(*out_shape, i, out); } // For the remaining shapes. - for (size_t i = params.num_out_data; i < g->outputs.size(); i++) { - uint32_t eid = idx.entry_id(g->outputs[i]); - (*out_shape)[i] = shapes[eid]; + for (size_t i = params.num_out_data; i < g.outputs.size(); i++) { + uint32_t eid = idx.entry_id(g.outputs[i]); + SHAPE_ASSIGN_CHECK(*out_shape, i, shapes[eid]); } - size_t num_states = g->outputs.size() - params.num_out_data; + size_t num_states = g.outputs.size() - params.num_out_data; for (size_t i = 0; i < num_states; i++) { size_t loc = params.in_state_locs[i]; CHECK((*out_shape)[i + params.num_out_data] == (*in_shape)[loc]); diff --git a/src/operator/nn/subgraph_op_common.cc b/src/operator/nn/subgraph_op_common.cc index 7efd48e0eed0..8221dbb7dfa6 100644 --- a/src/operator/nn/subgraph_op_common.cc +++ b/src/operator/nn/subgraph_op_common.cc @@ -42,11 +42,11 @@ bool InferSubgraphDataType(nnvm::Symbol &subgraph, std::vector *in_type, CHECK_EQ(input_nids.size(), in_type->size()); for (size_t i = 0; i < in_type->size(); i++) { auto eid = idx.entry_id(input_nids[i], 0); - (*in_type)[i] = dtypes[eid]; + TYPE_ASSIGN_CHECK(*in_type, i, dtypes[eid]); } for (size_t i = 0; i < g.outputs.size(); i++) - (*out_type)[i] = dtypes[idx.entry_id(g.outputs[i])]; + TYPE_ASSIGN_CHECK(*out_type, i, dtypes[idx.entry_id(g.outputs[i])]); return true; } @@ -73,14 +73,14 @@ bool InferSubgraphStorage(nnvm::Symbol &subgraph, CHECK_EQ(input_nids.size(), in_attrs->size()); for (size_t i = 0; i < in_attrs->size(); i++) { auto eid = idx.entry_id(input_nids[i], 0); - (*in_attrs)[i] = stypes[eid]; + STORAGE_TYPE_ASSIGN_CHECK(*in_attrs, i, stypes[eid]); } - *dispatch_mode = DispatchMode::kFComputeEx; + DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx); auto &outputs = idx.outputs(); CHECK(outputs.size() == out_attrs->size()); for (size_t i = 0; i < out_attrs->size(); i++) - (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; + STORAGE_TYPE_ASSIGN_CHECK(*out_attrs, i, stypes[idx.entry_id(outputs[i])]); return true; } @@ -142,11 +142,11 @@ bool InferSubgraphBackwardStorage(nnvm::Symbol &subgraph, std::move(storage_type_inputs), true); const auto& stypes = grad_graph.GetAttr("storage_type"); - *dispatch_mode = DispatchMode::kFComputeEx; + DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx); auto &outputs = idx.outputs(); CHECK(outputs.size() == out_attrs->size()); for (size_t i = 0; i < out_attrs->size(); i++) - (*out_attrs)[i] = stypes[idx.entry_id(outputs[i])]; + STORAGE_TYPE_ASSIGN_CHECK(*out_attrs, i, stypes[idx.entry_id(outputs[i])]); return true; } From f356460b28add16512497ee11717ff0e1f12990e Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 18 May 2018 22:27:54 +0000 Subject: [PATCH 48/71] fix lint. --- src/operator/nn/control_flow.cc | 11 +++-------- src/operator/nn/subgraph_op_common.cc | 7 ++++--- src/operator/nn/subgraph_op_common.h | 6 +++--- 3 files changed, 10 insertions(+), 14 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 8e242e950aff..caf77bda6029 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -90,9 +90,9 @@ class ForeachState { } }; -void ForeachState::Forward(std::vector cinputs, +void ForeachState::Forward(const std::vector &cinputs, const std::vector& req, - std::vector coutputs, bool is_recording) { + const std::vector &coutputs, bool is_recording) { using namespace nnvm; using namespace imperative; @@ -136,8 +136,6 @@ void ForeachState::Forward(std::vector cinputs, std::unordered_map > params; CachedOpPtr op = std::make_shared(subgraph_sym, kwargs, arg_names, params); - // TODO here we only changed the output arrays in the arguments. - // Will this be a problem? // TODO(zhengda) we need to avoid shape inference and memory plan whenever the op is // called. Currently, CachedOp allocates memory each time Forward is called. // I need to fix this once the PR for static memory allocation in CachedOp is @@ -145,7 +143,6 @@ void ForeachState::Forward(std::vector cinputs, op->Forward(nullptr, inputs, outputs); if (is_recording) { - // TODO does this have right inputs and outputs? all_outputs.push_back(coutputs); iter_ops.push_back(op); } @@ -186,8 +183,6 @@ void ForeachState::Backward(int iter_no, std::vector ograds, outputs.push_back(&igrads[i]); CHECK_EQ(outputs.size(), op->num_inputs()); - // TODO here we only changed the output arrays in the arguments. - // Will this be a problem? CHECK(!Imperative::AGInfo::IsNone(all_outputs[iter_no][0])); const nnvm::NodeEntry &node_entry = all_outputs[iter_no][0].GetAutogradEntry(); OpStatePtr state = Imperative::AGInfo::Get(node_entry.node).state; @@ -255,7 +250,7 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, std::vector *subg_out_prev = subg_outputs[(i + 1) % 2]; for (int j = 0; j < params.num_out_data; j++) (*subg_out_curr)[j] = outputs[j].At(i); - // When recording for backward computation, we should make sure + // When recording for backward computation, we should make sure // that output arrays are actually different in each iteration. if (ctx.need_grad && i < len - 1) { for (size_t j = params.num_out_data; j < subg_out_curr->size(); j++) diff --git a/src/operator/nn/subgraph_op_common.cc b/src/operator/nn/subgraph_op_common.cc index 8221dbb7dfa6..ac2218b062fe 100644 --- a/src/operator/nn/subgraph_op_common.cc +++ b/src/operator/nn/subgraph_op_common.cc @@ -24,7 +24,8 @@ namespace mxnet { namespace op { -bool InferSubgraphDataType(nnvm::Symbol &subgraph, std::vector *in_type, +bool InferSubgraphDataType(const nnvm::Symbol &subgraph, + std::vector *in_type, std::vector *out_type) { nnvm::DTypeVector dtype_inputs = *in_type; nnvm::Graph g; @@ -50,7 +51,7 @@ bool InferSubgraphDataType(nnvm::Symbol &subgraph, std::vector *in_type, return true; } -bool InferSubgraphStorage(nnvm::Symbol &subgraph, +bool InferSubgraphStorage(const nnvm::Symbol &subgraph, const int dev_mask, DispatchMode* dispatch_mode, std::vector *in_attrs, @@ -84,7 +85,7 @@ bool InferSubgraphStorage(nnvm::Symbol &subgraph, return true; } -bool InferSubgraphBackwardStorage(nnvm::Symbol &subgraph, +bool InferSubgraphBackwardStorage(const nnvm::Symbol &subgraph, const int dev_mask, DispatchMode* dispatch_mode, std::vector *in_attrs, diff --git a/src/operator/nn/subgraph_op_common.h b/src/operator/nn/subgraph_op_common.h index 6412bb45e8a3..1b6587953c78 100644 --- a/src/operator/nn/subgraph_op_common.h +++ b/src/operator/nn/subgraph_op_common.h @@ -27,14 +27,14 @@ namespace mxnet { namespace op { -bool InferSubgraphDataType(nnvm::Symbol &subgraph, std::vector *in_type, +bool InferSubgraphDataType(const nnvm::Symbol &subgraph, std::vector *in_type, std::vector *out_type); -bool InferSubgraphStorage(nnvm::Symbol &subgraph, +bool InferSubgraphStorage(const nnvm::Symbol &subgraph, const int dev_mask, DispatchMode* dispatch_mode, std::vector *in_attrs, std::vector *out_attrs); -bool InferSubgraphBackwardStorage(nnvm::Symbol &subgraph, +bool InferSubgraphBackwardStorage(const nnvm::Symbol &subgraph, const int dev_mask, DispatchMode* dispatch_mode, std::vector *in_attrs, From 6bc3d56ee3766a7ad322e2cc503295839cacf31e Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 18 May 2018 22:33:23 +0000 Subject: [PATCH 49/71] Fix lint. --- src/operator/nn/subgraph_op_common.h | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/operator/nn/subgraph_op_common.h b/src/operator/nn/subgraph_op_common.h index 1b6587953c78..2025a2baefa2 100644 --- a/src/operator/nn/subgraph_op_common.h +++ b/src/operator/nn/subgraph_op_common.h @@ -23,17 +23,32 @@ #include #include #include +#include namespace mxnet { namespace op { +/* + * Infer the data types of inputs and outputs of an operator that contains a + * subgraph. + */ bool InferSubgraphDataType(const nnvm::Symbol &subgraph, std::vector *in_type, std::vector *out_type); + +/* + * Infer the storage types of inputs and outputs of an operator that contains a + * subgraph. + */ bool InferSubgraphStorage(const nnvm::Symbol &subgraph, const int dev_mask, DispatchMode* dispatch_mode, std::vector *in_attrs, std::vector *out_attrs); + +/* + * Infer the storage types of inputs and outputs of the backward computation of + * an operator that contains a subgraph. + */ bool InferSubgraphBackwardStorage(const nnvm::Symbol &subgraph, const int dev_mask, DispatchMode* dispatch_mode, From 88ae9bff5bef177e08c2c210bab4b7b3db64b68f Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 19 May 2018 00:38:02 +0000 Subject: [PATCH 50/71] add doc. --- python/mxnet/symbol/contrib.py | 65 ++++++++++++++++++++++++++++++--- src/operator/nn/control_flow.cc | 2 +- 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index c236dffa63cd..bc5cd08c4583 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -99,11 +99,10 @@ def rand_zipfian(true_classes, num_sampled, range_max): expected_count_sampled = expected_prob_sampled * num_sampled return sampled_classes, expected_count_true, expected_count_sampled -def _get_graph_inputs(subg, name, prefix): +def _get_graph_inputs(subg): num_handles = ctypes.c_int(1000) handles = c_array(SymbolHandle, [SymbolHandle(0) for i in range(1000)]) - check_call(_LIB.MXSymbolGetInputSymbols(subg.handle, handles, - ctypes.byref(num_handles))) + check_call(_LIB.MXSymbolGetInputSymbols(subg.handle, handles, ctypes.byref(num_handles))) syms = [] for i in range(num_handles.value): @@ -111,7 +110,63 @@ def _get_graph_inputs(subg, name, prefix): syms.append(s) return syms -def foreach(func, data, init_states, back_prop=False, name="foreach"): +def foreach(func, data, init_states, name="foreach"): + """Run a for loop with user-defined computation over NDArrays on dimension 0. + + This operator simulates a for loop and func has the computation for an iteration + of the for loop. It runs the computation in func on each slice from the input + NDArrays. + + func takes two arguments as input and outputs a tuple of two elements, + as illustrated below: + + out, states = func(data1, states) + + data1 can be either a symbol or a list of symbols. If data is a symbol, + data1 is a symbol. Otherwise, data1 is a list of symbols and has the same + size as data. states is a list of symbols and have the same size as init_states. + Similarly, out can be either a symbol or a list of symbols, which are concatenated + as the first output of foreach; states from the last execution of func + are the second output of foreach. + + The computation done by this operator is equivalent to the pseudo code below + when the input data is NDArray: + + states = init_states + outs = [] + for i in data.shape[0]: + s = data[i] + out, states = func(s, states) + outs.append(out) + outs = stack(*outs) + + + Parameters + ---------- + func : a Python function. + Define computation in an iteration. + data: a symbol or a list of symbols. + The input data. + init_states: a list of symbols. + The initial values of the loop states. + name: string. + The name of the operator. + + Returns + ------- + outputs: a Symbol or a list of Symbols. + The output data concatenated from the output of all iterations. + states: a list of Symbols. + The loop states in the last iteration. + + Examples + -------- + >>> step = lambda data, states: (data + states[0], [states[0] * 2]) + >>> data = mx.sym.var('data') + >>> states = [mx.sym.var('state')] + >>> outs, states = mx.sym.contrib.foreach(step, data, states) + """ + assert isinstance(init_states, list), "init_states should be a list" states = [] @@ -148,7 +203,7 @@ def foreach(func, data, init_states, back_prop=False, name="foreach"): # all symbols will refer to different NDArrays. flat_out.append(symbol.op.identity(s)) g = symbol.Group(flat_out) - input_syms = _get_graph_inputs(g, name, "ro_var") + input_syms = _get_graph_inputs(g) if isinstance(data, list): num_inputs = len(data) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index caf77bda6029..3e8e733f572f 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -478,7 +478,7 @@ ForeachGradient(const nnvm::NodePtr& n, const std::vector& ogra } NNVM_REGISTER_OP(_foreach) -.describe(R"code(foreach)code" ADD_FILELINE) +.MXNET_DESCRIBE("Run a for loop over an NDArray with user-defined computation") .set_attr_parser(ParamParser) .set_attr("FInferStorageType", ForeachStorageType) .set_num_inputs([](const NodeAttrs& attrs) { From 1f03f015e5f46d3b295cf40a7d72b437e1527570 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 19 May 2018 00:39:39 +0000 Subject: [PATCH 51/71] undo modification in imperative.h --- include/mxnet/imperative.h | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/include/mxnet/imperative.h b/include/mxnet/imperative.h index c82c7afeb49b..758ce8513213 100644 --- a/include/mxnet/imperative.h +++ b/include/mxnet/imperative.h @@ -183,18 +183,18 @@ class Imperative { std::vector* p_save_inputs = nullptr, std::vector* p_save_outputs = nullptr); /*! \brief */ - static OpStatePtr Invoke(const Context& default_ctx, - const nnvm::NodeAttrs& attrs, - const std::vector& inputs, - const std::vector& outputs); + OpStatePtr Invoke(const Context& default_ctx, + const nnvm::NodeAttrs& attrs, + const std::vector& inputs, + const std::vector& outputs); /*! \brief */ - static OpStatePtr InvokeOp(const Context& ctx, - const nnvm::NodeAttrs& attrs, - const std::vector& inputs, - const std::vector& outputs, - const std::vector& req, - const DispatchMode dispatch_mode, - OpStatePtr state = OpStatePtr()); + OpStatePtr InvokeOp(const Context& ctx, + const nnvm::NodeAttrs& attrs, + const std::vector& inputs, + const std::vector& outputs, + const std::vector& req, + const DispatchMode dispatch_mode, + OpStatePtr state = OpStatePtr()); /*! \brief mark variables for computing gradients. */ void MarkVariables(const std::vector& variables, const std::vector& grad_reqs, From 28ba842a70f527fab58f700927f68cce7aa00761 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 19 May 2018 00:48:57 +0000 Subject: [PATCH 52/71] add doc and remove example code. --- python/mxnet/gluon/contrib/rnn/rnn_cell.py | 146 +-------------------- python/mxnet/ndarray/contrib.py | 62 ++++++++- 2 files changed, 60 insertions(+), 148 deletions(-) diff --git a/python/mxnet/gluon/contrib/rnn/rnn_cell.py b/python/mxnet/gluon/contrib/rnn/rnn_cell.py index dcb396a57613..1b9afee14bf2 100644 --- a/python/mxnet/gluon/contrib/rnn/rnn_cell.py +++ b/python/mxnet/gluon/contrib/rnn/rnn_cell.py @@ -17,15 +17,10 @@ # coding: utf-8 """Definition of various recurrent neural network cells.""" -__all__ = ['VariationalDropoutCell', 'LSTMPCell', 'SymHybridRNNCell', 'RNNCell'] +__all__ = ['VariationalDropoutCell', 'LSTMPCell'] -import inspect - -from .... import symbol, ndarray -from ....base import _as_list from ...rnn import BidirectionalCell, SequentialRNNCell, ModifierCell, HybridRecurrentCell from ...rnn.rnn_cell import _format_sequence, _get_begin_state, _mask_sequence_variable_length -from ...rnn.rnn_cell import RNNCell as GluonRNNCell from ... import tensor_types class VariationalDropoutCell(ModifierCell): @@ -320,142 +315,3 @@ def hybrid_forward(self, F, inputs, states, i2h_weight, return next_r, [next_r, next_c] # pylint: enable= arguments-differ - -class SymHybridRNNCell(HybridRecurrentCell): - def __init__(self, prefix=None, params=None): - super(SymHybridRNNCell, self).__init__(prefix=prefix, params=params) - - def unroll(self, inputs, begin_state=None, layout='TNC', - merge_outputs=None, valid_length=None): - # if this is a list, we can have unroll in the parent class to handle it. - if (isinstance(inputs, list)): - return super(SymHybridRNNCell, self).unroll(self, len(inputs), inputs, begin_state, - layout, merge_outputs, valid_length) - - self.reset() - batch_axis = layout.find('N') - axis = layout.find('T') - batch_size = 0 - if isinstance(inputs, symbol.Symbol): - F = symbol - else: - batch_size = inputs.shape[batch_axis] - F = ndarray - begin_state = _get_begin_state(self, F, begin_state, inputs, batch_size) - - states = begin_state - outputs = [] - all_states = [] - def iter_func(input, states): - return self(input, states) - outputs, last_states = F.contrib.foreach(iter_func, inputs, begin_state) - #if valid_length is not None: - # states = [F.SequenceLast(ele_list, - # sequence_length=valid_length, - # use_sequence_length=True, - # axis=0) - # for ele_list in all_states] - # outputs = F.SequenceMask(outputs, sequence_length=valid_length, use_sequence_length=True, - # axis=axis) - #outputs, _, _, _ = _format_sequence(length, outputs, layout, merge_outputs) - - return outputs, last_states - -class RNNCell(SymHybridRNNCell): - r"""Elman RNN recurrent neural network cell. - - Each call computes the following function: - - .. math:: - - h_t = \tanh(w_{ih} * x_t + b_{ih} + w_{hh} * h_{(t-1)} + b_{hh}) - - where :math:`h_t` is the hidden state at time `t`, and :math:`x_t` is the hidden - state of the previous layer at time `t` or :math:`input_t` for the first layer. - If nonlinearity='relu', then `ReLU` is used instead of `tanh`. - - Parameters - ---------- - hidden_size : int - Number of units in output symbol - activation : str or Symbol, default 'tanh' - Type of activation function. - i2h_weight_initializer : str or Initializer - Initializer for the input weights matrix, used for the linear - transformation of the inputs. - h2h_weight_initializer : str or Initializer - Initializer for the recurrent weights matrix, used for the linear - transformation of the recurrent state. - i2h_bias_initializer : str or Initializer - Initializer for the bias vector. - h2h_bias_initializer : str or Initializer - Initializer for the bias vector. - prefix : str, default 'rnn_' - Prefix for name of `Block`s - (and name of weight if params is `None`). - params : Parameter or None - Container for weight sharing between cells. - Created if `None`. - - - Inputs: - - **data**: input tensor with shape `(batch_size, input_size)`. - - **states**: a list of one initial recurrent state tensor with shape - `(batch_size, num_hidden)`. - - Outputs: - - **out**: output tensor with shape `(batch_size, num_hidden)`. - - **next_states**: a list of one output recurrent state tensor with the - same shape as `states`. - """ - def __init__(self, hidden_size, activation='tanh', - i2h_weight_initializer=None, h2h_weight_initializer=None, - i2h_bias_initializer='zeros', h2h_bias_initializer='zeros', - input_size=0, prefix=None, params=None): - super(RNNCell, self).__init__(prefix=prefix, params=params) - self._hidden_size = hidden_size - self._activation = activation - self._input_size = input_size - self.i2h_weight = self.params.get('i2h_weight', shape=(hidden_size, input_size), - init=i2h_weight_initializer, - allow_deferred_init=True) - self.h2h_weight = self.params.get('h2h_weight', shape=(hidden_size, hidden_size), - init=h2h_weight_initializer, - allow_deferred_init=True) - self.i2h_bias = self.params.get('i2h_bias', shape=(hidden_size,), - init=i2h_bias_initializer, - allow_deferred_init=True) - self.h2h_bias = self.params.get('h2h_bias', shape=(hidden_size,), - init=h2h_bias_initializer, - allow_deferred_init=True) - - def state_info(self, batch_size=0): - return [{'shape': (batch_size, self._hidden_size), '__layout__': 'NC'}] - - def _alias(self): - return 'rnn' - - def __repr__(self): - s = '{name}({mapping}' - if hasattr(self, '_activation'): - s += ', {_activation}' - s += ')' - shape = self.i2h_weight.shape - mapping = '{0} -> {1}'.format(shape[1] if shape[1] else None, shape[0]) - return s.format(name=self.__class__.__name__, - mapping=mapping, - **self.__dict__) - - def hybrid_forward(self, F, inputs, states, i2h_weight, - h2h_weight, i2h_bias, h2h_bias): - prefix = 't%d_'%self._counter - i2h = F.FullyConnected(data=inputs, weight=i2h_weight, bias=i2h_bias, - num_hidden=self._hidden_size, - name=prefix+'i2h') - h2h = F.FullyConnected(data=states[0], weight=h2h_weight, bias=h2h_bias, - num_hidden=self._hidden_size, - name=prefix+'h2h') - output = self._get_activation(F, i2h + h2h, self._activation, - name=prefix+'out') - - return output, [output] diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py index 7bf96d8f68f0..70709a0997e0 100644 --- a/python/mxnet/ndarray/contrib.py +++ b/python/mxnet/ndarray/contrib.py @@ -99,12 +99,68 @@ def rand_zipfian(true_classes, num_sampled, range_max, ctx=None): return sampled_classes, expected_count_true, expected_count_sampled # pylint: enable=line-too-long -def foreach(func, input, init_states, back_prop=False, name="foreach"): +def foreach(func, data, init_states, name="foreach"): + """Run a for loop with user-defined computation over NDArrays on dimension 0. + + This operator simulates a for loop and func has the computation for an iteration + of the for loop. It runs the computation in func on each slice from the input + NDArrays. + + func takes two arguments as input and outputs a tuple of two elements, + as illustrated below: + + out, states = func(data1, states) + + data1 can be either a symbol or a list of symbols. If data is a symbol, + data1 is a symbol. Otherwise, data1 is a list of symbols and has the same + size as data. states is a list of symbols and have the same size as init_states. + Similarly, out can be either a symbol or a list of symbols, which are concatenated + as the first output of foreach; states from the last execution of func + are the second output of foreach. + + The computation done by this operator is equivalent to the pseudo code below + when the input data is NDArray: + + states = init_states + outs = [] + for i in data.shape[0]: + s = data[i] + out, states = func(s, states) + outs.append(out) + outs = stack(*outs) + + + Parameters + ---------- + func : a Python function. + Define computation in an iteration. + data: a symbol or a list of symbols. + The input data. + init_states: a list of symbols. + The initial values of the loop states. + name: string. + The name of the operator. + + Returns + ------- + outputs: a Symbol or a list of Symbols. + The output data concatenated from the output of all iterations. + states: a list of Symbols. + The loop states in the last iteration. + + Examples + -------- + >>> step = lambda data, states: (data + states[0], [states[0] * 2]) + >>> data = mx.nd.random.uniform(shape=(2, 10)) + >>> states = [mx.nd.random.uniform(shape=(10))] + >>> outs, states = mx.nd.contrib.foreach(step, data, states) + """ + assert isinstance(init_states, list), "init_states should be a list" states = init_states outputs = [] - for i in range(input.shape[0]): - ele = input[i] + for i in range(data.shape[0]): + ele = data[i] outs, states = func(ele, states) outs = _as_list(outs) if (i == 0): From 14ca454912fbf974408b7283ea3bf18ea6aa5495 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 19 May 2018 00:59:44 +0000 Subject: [PATCH 53/71] fix lint. --- python/mxnet/symbol/contrib.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index bc5cd08c4583..8840f06a29e5 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -29,8 +29,8 @@ import ctypes from . import symbol -from ..base import _LIB, c_str, c_array, check_call -from ..base import SymbolHandle, NDArrayHandle, _as_list +from ..base import _LIB, c_array, check_call +from ..base import SymbolHandle, _as_list from ..attribute import AttrScope __all__ = ["rand_zipfian"] @@ -205,11 +205,6 @@ def foreach(func, data, init_states, name="foreach"): g = symbol.Group(flat_out) input_syms = _get_graph_inputs(g) - if isinstance(data, list): - num_inputs = len(data) - else: - num_inputs = 1 - # Here we need to find out how the input symbols are ordered as well as # where the loop states are located in the list of inputs. From 242455a21f4a82c455e71f2a8da86a09bbf34b0a Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 19 May 2018 01:01:55 +0000 Subject: [PATCH 54/71] fix lint. --- python/mxnet/ndarray/contrib.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py index 70709a0997e0..42fefcda3249 100644 --- a/python/mxnet/ndarray/contrib.py +++ b/python/mxnet/ndarray/contrib.py @@ -99,7 +99,7 @@ def rand_zipfian(true_classes, num_sampled, range_max, ctx=None): return sampled_classes, expected_count_true, expected_count_sampled # pylint: enable=line-too-long -def foreach(func, data, init_states, name="foreach"): +def foreach(func, data, init_states): """Run a for loop with user-defined computation over NDArrays on dimension 0. This operator simulates a for loop and func has the computation for an iteration From 3f5d207fc6c6924dd80ff57b8070fa0cc42c6908 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Sat, 19 May 2018 01:17:25 +0000 Subject: [PATCH 55/71] Fix lint. --- python/mxnet/ndarray/contrib.py | 14 +++++++------- python/mxnet/symbol/contrib.py | 4 ++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py index 42fefcda3249..77a705f4f98b 100644 --- a/python/mxnet/ndarray/contrib.py +++ b/python/mxnet/ndarray/contrib.py @@ -163,13 +163,13 @@ def foreach(func, data, init_states): ele = data[i] outs, states = func(ele, states) outs = _as_list(outs) - if (i == 0): + if i == 0: # outputs is a list of lists - for j in range(len(outs)): - outputs.append([outs[j]]) + for out in outs: + outputs.append([out]) else: - for j in range(len(outs)): - outputs[j].append(outs[j]) - for i in range(len(outputs)): - outputs[i] = stack(*outputs[i]) + for j, out in enumerate(outs): + outputs[j].append(out) + for out in outputs: + out = stack(*out) return (outputs, states) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index 8840f06a29e5..b1e448b4352a 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -19,6 +19,8 @@ # pylint: disable=wildcard-import, unused-wildcard-import """Contrib Symbol API of MXNet.""" import math +import ctypes + from .random import uniform from .symbol import Symbol try: @@ -26,8 +28,6 @@ except ImportError: pass -import ctypes - from . import symbol from ..base import _LIB, c_array, check_call from ..base import SymbolHandle, _as_list From cd7f94b42b2de4b137947f0d4fc0ab8440b5aa09 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 21 May 2018 20:20:07 +0000 Subject: [PATCH 56/71] make nd.foreach and sym.foreach consistent. --- python/mxnet/ndarray/contrib.py | 53 +++++++++++++++------- python/mxnet/symbol/contrib.py | 62 +++++++++++++++++--------- tests/python/unittest/test_operator.py | 48 ++++++++------------ 3 files changed, 99 insertions(+), 64 deletions(-) diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py index 77a705f4f98b..d6d72ad12150 100644 --- a/python/mxnet/ndarray/contrib.py +++ b/python/mxnet/ndarray/contrib.py @@ -22,7 +22,7 @@ from ..context import current_context from ..random import uniform from ..base import _as_list -from .op import stack +from . import ndarray try: from .gen_contrib import * except ImportError: @@ -111,10 +111,10 @@ def foreach(func, data, init_states): out, states = func(data1, states) - data1 can be either a symbol or a list of symbols. If data is a symbol, - data1 is a symbol. Otherwise, data1 is a list of symbols and has the same - size as data. states is a list of symbols and have the same size as init_states. - Similarly, out can be either a symbol or a list of symbols, which are concatenated + data1 can be either an NDArray or a list of NDArrays. If data is an NDArray, + data1 is an NDArray. Otherwise, data1 is a list of NDArrays and has the same + size as data. states is a list of NDArrays and have the same size as init_states. + Similarly, out can be either an NDArray or a list of NDArrays, which are concatenated as the first output of foreach; states from the last execution of func are the second output of foreach. @@ -134,18 +134,18 @@ def foreach(func, data, init_states): ---------- func : a Python function. Define computation in an iteration. - data: a symbol or a list of symbols. + data: an NDArray or a list of NDArrays. The input data. - init_states: a list of symbols. + init_states: an NDArray or a list of NDArrays. The initial values of the loop states. name: string. The name of the operator. Returns ------- - outputs: a Symbol or a list of Symbols. + outputs: an NDArray or a list of NDArrays. The output data concatenated from the output of all iterations. - states: a list of Symbols. + states: a list of NDArrays. The loop states in the last iteration. Examples @@ -156,12 +156,32 @@ def foreach(func, data, init_states): >>> outs, states = mx.nd.contrib.foreach(step, data, states) """ - assert isinstance(init_states, list), "init_states should be a list" + def check_input(inputs, in_type, msg): + is_NDArray_or_list = True + if isinstance(inputs, list): + for i in inputs: + if not isinstance(i, in_type): + is_NDArray_or_list = False + break + else: + is_NDArray_or_list = isinstance(inputs, in_type) + assert is_NDArray_or_list, msg + + check_input(data, ndarray.NDArray, "data should be an NDArray or a list of NDArrays") + check_input(init_states, ndarray.NDArray, + "init_states should be an NDArray or a list of NDArrays") + + not_data_list = isinstance(data, ndarray.NDArray) + not_state_list = isinstance(init_states, ndarray.NDArray) + num_iters = data.shape[0] if not_data_list else data[0].shape[0] states = init_states outputs = [] - for i in range(data.shape[0]): - ele = data[i] - outs, states = func(ele, states) + for i in range(num_iters): + if not_data_list: + eles = data[i] + else: + eles = [d[i] for d in data] + outs, states = func(eles, states) outs = _as_list(outs) if i == 0: # outputs is a list of lists @@ -170,6 +190,9 @@ def foreach(func, data, init_states): else: for j, out in enumerate(outs): outputs[j].append(out) - for out in outputs: - out = stack(*out) + for j, out in enumerate(outputs): + outputs[j] = ndarray.op.stack(*out) + + if not_data_list: + outputs = outputs[0] return (outputs, states) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index b1e448b4352a..5f13e365a0cc 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -147,7 +147,7 @@ def foreach(func, data, init_states, name="foreach"): Define computation in an iteration. data: a symbol or a list of symbols. The input data. - init_states: a list of symbols. + init_states: a symbol or a list of symbols. The initial values of the loop states. name: string. The name of the operator. @@ -167,8 +167,21 @@ def foreach(func, data, init_states, name="foreach"): >>> outs, states = mx.sym.contrib.foreach(step, data, states) """ - assert isinstance(init_states, list), "init_states should be a list" - states = [] + def check_data(inputs, in_type, msg): + is_NDArray_or_list = True + if isinstance(inputs, list): + for i in inputs: + if not isinstance(i, in_type): + is_NDArray_or_list = False + break + else: + is_NDArray_or_list = isinstance(inputs, in_type) + assert is_NDArray_or_list, msg + + check_data(data, symbol.Symbol, "data should be an NDArray or a list of NDArrays") + check_data(init_states, symbol.Symbol, + "init_states should be an NDArray or a list of NDArrays") + not_state_list = isinstance(init_states, symbol.Symbol) # TODO(zhengda) If the input python function references to the symbols outside # the python function, we need to prune the computation graph constructed from @@ -179,29 +192,33 @@ def foreach(func, data, init_states, name="foreach"): in_eles = [symbol.var(sym.name) for sym in data] else: in_eles = symbol.var(data.name) - for s in init_states: - states.append(symbol.var(s.name)) - - sym_out = func(in_eles, states) - # The function should return a tuple. The first element goes to - # the output of the function. The second element is a list. - assert isinstance(sym_out, tuple), "func should return a tuple (out, states)" - assert isinstance(sym_out[1], list), \ - "the second element in the returned tuple should be a list" - assert len(sym_out[1]) == len(init_states), \ + if isinstance(init_states, list): + states = [symbol.var(s.name) for s in init_states] + else: + states = symbol.var(init_states.name) + sym_out, sym_states = func(in_eles, states) + + check_data(sym_out, symbol.Symbol, "the output should be an NDArray or a list of NDArrays") + check_data(sym_states, symbol.Symbol, + "the output states should be an NDArray or a list of NDArrays") + if isinstance(sym_states, list): + assert isinstance(init_states, list) and len(sym_states) == len(init_states), \ "the number of output states (%d) should be the same as input states (%d)" \ - % (len(sym_out[1]), len(init_states)) + % (len(sym_states), len(init_states)) - if isinstance(sym_out[0], list): - flat_out = sym_out[0] - else: - flat_out = [sym_out[0]] - num_out_data = len(flat_out) - for s in sym_out[1]: + if isinstance(sym_out, list): + flat_out = sym_out + else: + flat_out = [sym_out] + num_out_data = len(flat_out) + if isinstance(sym_states, list): + for s in sym_states: # There is a problem if the outputs are the same as the inputs # or the first output. By calling identity, we can make sure that # all symbols will refer to different NDArrays. flat_out.append(symbol.op.identity(s)) + else: + flat_out.append(symbol.op.identity(sym_states)) g = symbol.Group(flat_out) input_syms = _get_graph_inputs(g) @@ -248,4 +265,9 @@ def foreach(func, data, init_states, name="foreach"): for i in range(num_states): states.append(ret[num_outputs - num_states + i]) + if not_state_list: + # If there is only one input state, there should be only one output state. + assert len(states) == 1 + states = states[0] + return (outs, states) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 3ed44c80df04..4fce81a7ae14 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5685,13 +5685,12 @@ def step3(in1, states, free): def verify_foreach(step, in_syms, state_syms, free_syms, in_arrs, init_states, frees, out_grads, is_train=True): step_sym = lambda in_syms, state_syms : step(in_syms, state_syms, free_syms) - step_imp = lambda in_arrs, state_arrs : step(in_arrs, state_arrs, frees) - out = mx.sym.contrib.foreach(step_sym, in_syms, state_syms) - out1 = _as_list(out[0]) - for i in range(len(out1)): - out1[i] = out1[i] * 2 - out1.extend(out[1]) - out = mx.sym.Group(out1) + res, states = mx.sym.contrib.foreach(step_sym, in_syms, state_syms) + out = _as_list(res) + for i in range(len(out)): + out[i] = out[i] * 2 + out.extend(states) + out = mx.sym.Group(out) arr_grads = [] arg_dict = {} arg_grad_dict = {} @@ -5720,8 +5719,7 @@ def verify_foreach(step, in_syms, state_syms, free_syms, name = name[1:] gin_order.append(int(name)) - e = out.bind(ctx=mx.cpu(), args=arg_dict, args_grad=arg_grad_dict, - ) + e = out.bind(ctx=mx.cpu(), args=arg_dict, args_grad=arg_grad_dict) e.forward(is_train=is_train) if (is_train): # backward @@ -5739,29 +5737,21 @@ def verify_foreach(step, in_syms, state_syms, free_syms, arr.attach_grad() for arr in frees: arr.attach_grad() + step_imp = lambda in_arrs, state_arrs : step(in_arrs, state_arrs, frees) with mx.autograd.record(): states = [mx.nd.expand_dims(s, 0) for s in init_states] - if isinstance(in_arrs, list): - num_iters = in_arrs[0].shape[0] + res, states = mx.nd.contrib.foreach(step_imp, in_arrs, init_states) + + res2 = _as_list(res) + for i in range(len(res2)): + res2[i] = res2[i] * 2 + if isinstance(states, list): + states = [mx.nd.expand_dims(s, 0) for s in states] + res2.extend(states) else: - num_iters = in_arrs.shape[0] - - for i in range(num_iters): - if isinstance(in_arrs, list): - data = [mx.nd.expand_dims(arr[i], 0) for arr in in_arrs] - else: - data = mx.nd.expand_dims(in_arrs[i], 0) - tmp_res = step_imp(data, states) - tmp_res1 = _as_list(tmp_res[0]) - for i in range(len(tmp_res1)): - res[i].append(tmp_res1[i]) - states = tmp_res[1] - res2 = [] - for l in res: - res2.append(mx.nd.concat(*l, dim=0) * 2) - tmp_res2 = res2[:] - tmp_res2.extend(tmp_res[1]) - res = mx.nd.concat(*tmp_res2, dim=0) + states = mx.nd.expand_dims(states, 0) + res2.append(states) + res = mx.nd.concat(*res2, dim=0) tmp_grads = out_grads[0][:] tmp_grads1 = [mx.nd.expand_dims(grad, 0) for grad in out_grads[1]] From 3936affd9d0994170a16b496581484c13e8e5777 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 21 May 2018 20:28:54 +0000 Subject: [PATCH 57/71] fix compile error. --- src/operator/nn/control_flow.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 3e8e733f572f..76f33f9ff8f8 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -90,9 +90,9 @@ class ForeachState { } }; -void ForeachState::Forward(const std::vector &cinputs, +void ForeachState::Forward(std::vector cinputs, const std::vector& req, - const std::vector &coutputs, bool is_recording) { + std::vector coutputs, bool is_recording) { using namespace nnvm; using namespace imperative; From fb23c9073b687afe1542eba811344fede03cac59 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 16 May 2018 22:57:54 +0000 Subject: [PATCH 58/71] Fix bugs in MKLDNN. --- src/ndarray/ndarray.cc | 60 ++++++++++++++++++++++++++---------------- 1 file changed, 38 insertions(+), 22 deletions(-) diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc index 1accd6f9f1cb..764711f020ff 100644 --- a/src/ndarray/ndarray.cc +++ b/src/ndarray/ndarray.cc @@ -200,6 +200,7 @@ NDArray NDArray::MKLDNNDataReshape(const TShape &shape) const { ret.ptr_->delay_alloc = false; ret.ptr_->static_data = true; ret.byte_offset_ = byte_offset_; + ret.reuse_ = false; return ret; } } @@ -217,6 +218,7 @@ NDArray NDArray::Reshape(const TShape &shape) const { // Otherwise, reshape only works on the default layout. CHECK_EQ(storage_type(), kDefaultStorage); ret.shape_ = shape; + ret.reuse_ = false; return ret; } @@ -249,6 +251,7 @@ NDArray NDArray::Slice(index_t begin, index_t end) const { MSHADOW_TYPE_SWITCH(ret.dtype(), DType, { ret.byte_offset_ += begin * length * sizeof(DType); }); + ret.reuse_ = false; ret.shape_[0] = end - begin; return ret; } @@ -555,6 +558,7 @@ NDArray NDArray::Reorder2Default() const { // reshape as needed ret.shape_ = shape_; ret.byte_offset_ = byte_offset_; + ret.reuse_ = false; return ret; } @@ -584,39 +588,39 @@ void NDArray::MKLDNNDataReorderAsync(const mkldnn::memory::primitive_desc &desc) const mkldnn::memory *NDArray::GetMKLDNNData() const { CHECK(storage_type() == kDefaultStorage); + bool is_view = IsView(); if (IsMKLDNNData()) { // If this array uses MKLDNN layout, we have to make sure it's not a view. // Otherwise, we'll have to change the layout inside the array. - CHECK(!IsView()); + CHECK(!is_view); MKLDNNStream::Get()->RegisterMem(ptr_->mkl_mem_->GetMem()); // If this array uses MKLDNN format, we should return now. Otherwise, // SetMKLMem may mess up mkl_mem_. return ptr_->mkl_mem_->GetRaw(); - } - ptr_->SetMKLMem(IsView() ? ptr_->storage_shape : shape_, dtype_); - MKLDNNStream::Get()->RegisterMem(ptr_->mkl_mem_->GetMem()); - if (IsView()) { - mkldnn::memory::primitive_desc pd = ptr_->mkl_mem_->GetPrimitiveDesc(); - // Sliced array must use the default layout. - CHECK_EQ(GetDefaultFormat(pd.desc()), pd.desc().data.format); - void *off_addr = static_cast(ptr_->mkl_mem_->GetDataHandle()) - + byte_offset_; - + } else if (is_view) { + // If this is a view, we can't create a MKLDNN memory for the chunk + // because we don't have the complete data type and shape information for + // the chunk. + void *off_addr = static_cast(ptr_->shandle.dptr) + byte_offset_; // Create the primitive desc for the new mkldnn memory. mkldnn::memory::dims dims(shape().ndim()); for (size_t i = 0; i < dims.size(); i++) dims[i] = shape()[i]; mkldnn::memory::format cpp_format = static_cast( GetDefaultFormat(shape().ndim())); - mkldnn::memory::data_type cpp_type = static_cast( - pd.desc().data.data_type); + mkldnn::memory::data_type cpp_type = get_mkldnn_type(dtype_); mkldnn::memory::desc data_md(dims, cpp_type, cpp_format); - mkldnn::memory::primitive_desc new_pd(data_md, pd.get_engine()); + mkldnn::memory::primitive_desc new_pd(data_md, + CpuEngine::Get()->get_engine()); std::shared_ptr ret(new mkldnn::memory(new_pd, off_addr)); MKLDNNStream::Get()->RegisterMem(ret); return ret.get(); } else { + // If this isn't a view, we can create a MKLDNN memory and store it in the + // chunk. + ptr_->SetMKLMem(shape_, dtype_); + MKLDNNStream::Get()->RegisterMem(ptr_->mkl_mem_->GetMem()); return ptr_->mkl_mem_->GetRaw(); } } @@ -637,10 +641,9 @@ void NDArray::CopyFrom(const mkldnn::memory &mem) { MKLDNNStream *stream = MKLDNNStream::Get(); // If this array uses MKLDNN layout, we have to make sure it's not a view. // Otherwise, we'll have to change the layout inside the array. - if (IsMKLDNNData()) - CHECK(!IsView()); - ptr_->SetMKLMem(IsView() ? ptr_->storage_shape : shape_, - dtype_); + + CHECK(!IsView()); + ptr_->SetMKLMem(shape_, dtype_); stream->RegisterMem(ptr_->mkl_mem_->GetMem()); mkldnn::memory::desc from_desc = mem.get_primitive_desc().desc(); mkldnn::memory::desc this_desc = ptr_->mkl_mem_->GetPrimitiveDesc().desc(); @@ -713,9 +716,6 @@ mkldnn::memory::primitive_desc GetPrimitiveDesc(mkldnn::memory::primitive_desc p mkldnn_memory_format_t format); mkldnn::memory *NDArray::CreateMKLDNNData(const mkldnn::memory::primitive_desc &desc) { - // This array shouldn't be a view. - CHECK(!IsView()); - if (desc.get_size() != shape().Size() * GetTypeSize(dtype_)) { LOG(FATAL) << "The size of NDArray doesn't match the requested MKLDNN memory desc"; return nullptr; @@ -726,10 +726,26 @@ mkldnn::memory *NDArray::CreateMKLDNNData(const mkldnn::memory::primitive_desc & mkldnn_memory_format_t def_format = GetDefaultFormat(_desc.desc()); // If the required format is a default format, we don't need to worry about the shape. // If the shape isn't the same, it actually implicitly reshapes data. - if (required_format == def_format) { + if (required_format == def_format && !IsView()) { ptr_->SetMKLMem(shape_, dtype_); MKLDNNStream::Get()->RegisterMem(ptr_->mkl_mem_->GetMem()); return GetMKLDNNExact(ptr_->mkl_mem_->GetRaw(), desc); + } else if (required_format == def_format) { + ptr_->CheckAndAlloc(); + CHECK(ptr_->shandle.dptr); + // When this is a view and a user wants the default layout, we can simply + // create a new mkldnn memory that points to the right memory. + std::shared_ptr mem(new mkldnn::memory( + desc, ptr_->shandle.dptr + byte_offset_)); + MKLDNNStream::Get()->RegisterMem(mem); + return mem.get(); + } else if (IsView()) { + // If this is a view and a user wants to write data to it with special + // a MKLDNN format, we should reorder the data in the array and return NULL. + // In this way, the user will create a new NDArray for the special format + // and copy data back. + ptr_->Reorder2Default(); + return nullptr; } if (ptr_->mkl_mem_) From 0a27af1b4dfdb672e3b735a415e63432e80c053b Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 21 May 2018 22:40:23 +0000 Subject: [PATCH 59/71] address comments. --- include/mxnet/c_api.h | 4 ++-- include/mxnet/ndarray.h | 2 +- python/mxnet/ndarray/contrib.py | 27 +++++++++++---------------- python/mxnet/symbol/contrib.py | 18 +++++++++--------- src/c_api/c_api_symbolic.cc | 10 +++++----- src/executor/attach_op_execs_pass.cc | 3 ++- src/imperative/imperative_utils.h | 7 +++---- src/operator/nn/control_flow.cc | 14 +++++++++----- 8 files changed, 42 insertions(+), 43 deletions(-) diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h index 4c37ba9d5400..a8e87271e166 100644 --- a/include/mxnet/c_api.h +++ b/include/mxnet/c_api.h @@ -1063,8 +1063,8 @@ MXNET_DLL int MXSymbolGetAtomicSymbolName(AtomicSymbolCreator creator, * \param outs The input symbols of the graph. * \param out_size the number of input symbols returned. */ -MXNET_DLL int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **outs, - int *out_size); +MXNET_DLL int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **inputs, + int *input_size); /*! * \brief Get the detailed information about atomic symbol. diff --git a/include/mxnet/ndarray.h b/include/mxnet/ndarray.h index c6f181c120d2..897b5e882aaa 100644 --- a/include/mxnet/ndarray.h +++ b/include/mxnet/ndarray.h @@ -693,7 +693,7 @@ class NDArray { NDArray MKLDNNDataReshape(const TShape &shape) const; #endif - const nnvm::NodeEntry &GetAutogradEntry() const { + const nnvm::NodeEntry &entry() const { return entry_; } diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py index d6d72ad12150..d68698f71d66 100644 --- a/python/mxnet/ndarray/contrib.py +++ b/python/mxnet/ndarray/contrib.py @@ -99,23 +99,23 @@ def rand_zipfian(true_classes, num_sampled, range_max, ctx=None): return sampled_classes, expected_count_true, expected_count_sampled # pylint: enable=line-too-long -def foreach(func, data, init_states): +def foreach(body, data, init_states): """Run a for loop with user-defined computation over NDArrays on dimension 0. - This operator simulates a for loop and func has the computation for an iteration - of the for loop. It runs the computation in func on each slice from the input + This operator simulates a for loop and body has the computation for an iteration + of the for loop. It runs the computation in body on each slice from the input NDArrays. - func takes two arguments as input and outputs a tuple of two elements, + body takes two arguments as input and outputs a tuple of two elements, as illustrated below: - out, states = func(data1, states) + out, states = body(data1, states) data1 can be either an NDArray or a list of NDArrays. If data is an NDArray, data1 is an NDArray. Otherwise, data1 is a list of NDArrays and has the same size as data. states is a list of NDArrays and have the same size as init_states. Similarly, out can be either an NDArray or a list of NDArrays, which are concatenated - as the first output of foreach; states from the last execution of func + as the first output of foreach; states from the last execution of body are the second output of foreach. The computation done by this operator is equivalent to the pseudo code below @@ -125,14 +125,14 @@ def foreach(func, data, init_states): outs = [] for i in data.shape[0]: s = data[i] - out, states = func(s, states) + out, states = body(s, states) outs.append(out) outs = stack(*outs) Parameters ---------- - func : a Python function. + body : a Python function. Define computation in an iteration. data: an NDArray or a list of NDArrays. The input data. @@ -181,15 +181,10 @@ def check_input(inputs, in_type, msg): eles = data[i] else: eles = [d[i] for d in data] - outs, states = func(eles, states) + outs, states = body(eles, states) outs = _as_list(outs) - if i == 0: - # outputs is a list of lists - for out in outs: - outputs.append([out]) - else: - for j, out in enumerate(outs): - outputs[j].append(out) + outputs.append(outs) + outputs = zip(*outputs) for j, out in enumerate(outputs): outputs[j] = ndarray.op.stack(*out) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index 5f13e365a0cc..383582393ac7 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -110,23 +110,23 @@ def _get_graph_inputs(subg): syms.append(s) return syms -def foreach(func, data, init_states, name="foreach"): +def foreach(body, data, init_states, name="foreach"): """Run a for loop with user-defined computation over NDArrays on dimension 0. - This operator simulates a for loop and func has the computation for an iteration - of the for loop. It runs the computation in func on each slice from the input + This operator simulates a for loop and body has the computation for an iteration + of the for loop. It runs the computation in body on each slice from the input NDArrays. - func takes two arguments as input and outputs a tuple of two elements, + body takes two arguments as input and outputs a tuple of two elements, as illustrated below: - out, states = func(data1, states) + out, states = body(data1, states) data1 can be either a symbol or a list of symbols. If data is a symbol, data1 is a symbol. Otherwise, data1 is a list of symbols and has the same size as data. states is a list of symbols and have the same size as init_states. Similarly, out can be either a symbol or a list of symbols, which are concatenated - as the first output of foreach; states from the last execution of func + as the first output of foreach; states from the last execution of body are the second output of foreach. The computation done by this operator is equivalent to the pseudo code below @@ -136,14 +136,14 @@ def foreach(func, data, init_states, name="foreach"): outs = [] for i in data.shape[0]: s = data[i] - out, states = func(s, states) + out, states = body(s, states) outs.append(out) outs = stack(*outs) Parameters ---------- - func : a Python function. + body : a Python function. Define computation in an iteration. data: a symbol or a list of symbols. The input data. @@ -196,7 +196,7 @@ def check_data(inputs, in_type, msg): states = [symbol.var(s.name) for s in init_states] else: states = symbol.var(init_states.name) - sym_out, sym_states = func(in_eles, states) + sym_out, sym_states = body(in_eles, states) check_data(sym_out, symbol.Symbol, "the output should be an NDArray or a list of NDArrays") check_data(sym_states, symbol.Symbol, diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index 3ffa07256266..34819e83d506 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -345,14 +345,14 @@ int MXSymbolGetAtomicSymbolName(AtomicSymbolCreator creator, API_END(); } -int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **out_arr, int *out_size) { +int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **input_arr, int *input_size) { API_BEGIN(); nnvm::Symbol *s = static_cast(sym); nnvm::Graph g; g.outputs = s->outputs; std::vector input_syms; const nnvm::IndexedGraph& idx = g.indexed_graph(); - size_t max_out_size = *out_size; + size_t max_input_size = *input_size; // Go through all nodes and return the ones representing variables. for (size_t i = 0; i < idx.num_nodes(); i++) { const nnvm::Node &n = *idx[i].source; @@ -365,9 +365,9 @@ int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **out_arr, int *out_s } } } - CHECK(input_syms.size() <= max_out_size); - *out_size = input_syms.size(); - memcpy(out_arr, input_syms.data(), sizeof(*out_arr) * input_syms.size()); + CHECK(input_syms.size() <= max_input_size); + *input_size = input_syms.size(); + memcpy(input_arr, input_syms.data(), sizeof(*input_arr) * input_syms.size()); API_END_HANDLE_ERROR(); } diff --git a/src/executor/attach_op_execs_pass.cc b/src/executor/attach_op_execs_pass.cc index 06b01b4f67e3..b90aa83099ae 100644 --- a/src/executor/attach_op_execs_pass.cc +++ b/src/executor/attach_op_execs_pass.cc @@ -138,7 +138,8 @@ class StatefulComputeExecutor : public StorageFallbackOpExecutor { return state_.get_var(); } - explicit StatefulComputeExecutor(const NodeAttrs& attrs, const OpStatePtr& state, + explicit StatefulComputeExecutor(const NodeAttrs& attrs, + const OpStatePtr& state, const FStatefulCompute& fcompute, ExecType exec_type, const std::vector &mutate_idx) diff --git a/src/imperative/imperative_utils.h b/src/imperative/imperative_utils.h index 0051becc38f0..1135c0d2d416 100644 --- a/src/imperative/imperative_utils.h +++ b/src/imperative/imperative_utils.h @@ -359,6 +359,7 @@ inline void PushFCompute(const FCompute& fn, static auto& fexec_type = nnvm::Op::GetAttr("FExecType"); bool is_train = Imperative::Get()->is_training(); + bool need_grad = Imperative::Get()->is_recording(); ExecType exec_type = fexec_type.count(op) ? fexec_type[op](attrs) : ExecType::kSync; CHECK(exec_type == ExecType::kSync); std::vector inputs, outputs; @@ -379,7 +380,6 @@ inline void PushFCompute(const FCompute& fn, &input_blobs, &output_blobs, &pre_temp_src, &pre_temp_dst, &post_temp_src, &post_temp_dst, &in_temp_idx_map, mutate_idx); // setup context - bool need_grad = Imperative::Get()->is_recording(); OpContext opctx{need_grad, is_train, rctx, engine::CallbackOnComplete(), requested}; bool is_gpu = ctx.dev_mask() == gpu::kDevMask; // pre-fcompute fallback, cast to default storage type @@ -407,11 +407,11 @@ inline void PushFComputeEx(const FComputeEx& fn, static auto& fexec_type = nnvm::Op::GetAttr("FExecType"); bool is_train = Imperative::Get()->is_training(); + bool need_grad = Imperative::Get()->is_recording(); ExecType exec_type = fexec_type.count(op) ? fexec_type[op](attrs) : ExecType::kSync; std::vector inputs, outputs; DerefInputOutput(p_inputs, p_outputs, &inputs, &outputs); const auto& run = [=](RunContext rctx) { - bool need_grad = Imperative::Get()->is_recording(); OpContext opctx{need_grad, is_train, rctx, engine::CallbackOnComplete(), requested}; #if MXNET_USE_MKLDNN == 1 InvalidateOutputs(outputs, req); @@ -447,6 +447,7 @@ inline void PushOperator(const OpStatePtr& state, static auto& fexec_type = nnvm::Op::GetAttr("FExecType"); bool is_train = Imperative::Get()->is_training(); + bool need_grad = Imperative::Get()->is_recording(); ExecType exec_type = fexec_type.count(op) ? fexec_type[op](attrs) : ExecType::kSync; std::vector inputs, outputs; DerefInputOutput(p_inputs, p_outputs, &inputs, &outputs); @@ -458,7 +459,6 @@ inline void PushOperator(const OpStatePtr& state, if (fcompute_ex != nullptr && dispatch_mode == DispatchMode::kFComputeEx) { const auto& run = [=](RunContext rctx, engine::CallbackOnComplete on_complete) { - bool need_grad = Imperative::Get()->is_recording(); OpContext opctx{need_grad, is_train, rctx, on_complete, requested}; #if MXNET_USE_MKLDNN == 1 InvalidateOutputs(outputs, req); @@ -492,7 +492,6 @@ inline void PushOperator(const OpStatePtr& state, << "for stateful operator " << op->name; const auto& run = [=](RunContext rctx, engine::CallbackOnComplete on_complete) { - bool need_grad = Imperative::Get()->is_recording(); OpContext opctx{need_grad, is_train, rctx, on_complete, requested}; std::vector input_blobs, output_blobs; diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 76f33f9ff8f8..6cb28de02afa 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -78,8 +78,10 @@ class ForeachState { void Forward(std::vector cinputs, const std::vector& req, - std::vector coutputs, bool is_recording); - void Backward(int iter_no, std::vector ograds, + std::vector coutputs, + bool is_recording); + void Backward(int iter_no, + std::vector ograds, const std::vector &req, std::vector igrads); void Cleanup() { @@ -92,7 +94,8 @@ class ForeachState { void ForeachState::Forward(std::vector cinputs, const std::vector& req, - std::vector coutputs, bool is_recording) { + std::vector coutputs, + bool is_recording) { using namespace nnvm; using namespace imperative; @@ -150,7 +153,8 @@ void ForeachState::Forward(std::vector cinputs, Imperative::Get()->set_is_recording(orig_is_record); } -void ForeachState::Backward(int iter_no, std::vector ograds, +void ForeachState::Backward(int iter_no, + std::vector ograds, const std::vector &req, std::vector igrads) { using namespace nnvm; @@ -184,7 +188,7 @@ void ForeachState::Backward(int iter_no, std::vector ograds, CHECK_EQ(outputs.size(), op->num_inputs()); CHECK(!Imperative::AGInfo::IsNone(all_outputs[iter_no][0])); - const nnvm::NodeEntry &node_entry = all_outputs[iter_no][0].GetAutogradEntry(); + const nnvm::NodeEntry &node_entry = all_outputs[iter_no][0].entry(); OpStatePtr state = Imperative::AGInfo::Get(node_entry.node).state; op->Backward(false, state, inputs, req, outputs); } From 76006ddbdbf45b0df27a0cba6975ffd525565bb2 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Mon, 21 May 2018 23:08:40 +0000 Subject: [PATCH 60/71] update. --- python/mxnet/symbol/contrib.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index 383582393ac7..2d067555e8dc 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -111,7 +111,7 @@ def _get_graph_inputs(subg): return syms def foreach(body, data, init_states, name="foreach"): - """Run a for loop with user-defined computation over NDArrays on dimension 0. + """Run a for loop with user-defined computation over Symbols on dimension 0. This operator simulates a for loop and body has the computation for an iteration of the for loop. It runs the computation in body on each slice from the input From a42a5a0aea5f08ebc747b5f7156d7419c89973bb Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 22 May 2018 00:31:45 +0000 Subject: [PATCH 61/71] check for loop only works for dense arrays. --- src/operator/nn/control_flow.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/operator/nn/control_flow.cc b/src/operator/nn/control_flow.cc index 6cb28de02afa..df28d5cb2cae 100644 --- a/src/operator/nn/control_flow.cc +++ b/src/operator/nn/control_flow.cc @@ -211,6 +211,9 @@ static void ForeachComputeExCPU(const OpStatePtr& state_ptr, } for (size_t i = 0; i < (size_t) params.num_out_data; i++) CHECK_EQ(len, outputs[i].shape()[iter_dim]); + for (const auto &arr : outputs) + CHECK_EQ(arr.storage_type(), kDefaultStorage) + << "The for operator doesn't support the sparse format"; // Initialize the outputs of the subgraph is a little trickier. // The states from the previous iteration are used as the inputs of the next @@ -302,6 +305,9 @@ static void ForeachGradComputeExCPU(const OpStatePtr& state_ptr, const ForeachParam& params = state.params; CHECK_EQ(outputs.size(), (size_t) params.num_args - 1); CHECK_GT(params.in_data_locs.ndim(), 0); + for (const auto &arr : outputs) + CHECK_EQ(arr.storage_type(), kDefaultStorage) + << "The for operator doesn't support the sparse format"; size_t iter_dim = 0; std::unordered_set in_data_locs(params.in_data_locs.begin(), params.in_data_locs.end()); From d2eb153b10f2eac1405ed4cc95fa7ef310fc60c6 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 22 May 2018 00:32:30 +0000 Subject: [PATCH 62/71] move control flow op out of nn/ --- src/operator/{nn => }/control_flow.cc | 0 src/operator/{nn => }/subgraph_op_common.cc | 0 src/operator/{nn => }/subgraph_op_common.h | 6 +++--- 3 files changed, 3 insertions(+), 3 deletions(-) rename src/operator/{nn => }/control_flow.cc (100%) rename src/operator/{nn => }/subgraph_op_common.cc (100%) rename src/operator/{nn => }/subgraph_op_common.h (93%) diff --git a/src/operator/nn/control_flow.cc b/src/operator/control_flow.cc similarity index 100% rename from src/operator/nn/control_flow.cc rename to src/operator/control_flow.cc diff --git a/src/operator/nn/subgraph_op_common.cc b/src/operator/subgraph_op_common.cc similarity index 100% rename from src/operator/nn/subgraph_op_common.cc rename to src/operator/subgraph_op_common.cc diff --git a/src/operator/nn/subgraph_op_common.h b/src/operator/subgraph_op_common.h similarity index 93% rename from src/operator/nn/subgraph_op_common.h rename to src/operator/subgraph_op_common.h index 2025a2baefa2..25cbd60f5b63 100644 --- a/src/operator/nn/subgraph_op_common.h +++ b/src/operator/subgraph_op_common.h @@ -17,8 +17,8 @@ * under the License. */ -#ifndef MXNET_OPERATOR_NN_SUBGRAPH_OP_COMMON_H_ -#define MXNET_OPERATOR_NN_SUBGRAPH_OP_COMMON_H_ +#ifndef MXNET_OPERATOR_SUBGRAPH_OP_COMMON_H_ +#define MXNET_OPERATOR_SUBGRAPH_OP_COMMON_H_ #include #include @@ -58,4 +58,4 @@ bool InferSubgraphBackwardStorage(const nnvm::Symbol &subgraph, } // namespace op } // namespace mxnet -#endif // MXNET_OPERATOR_NN_SUBGRAPH_OP_COMMON_H_ +#endif // MXNET_OPERATOR_SUBGRAPH_OP_COMMON_H_ From 767857fc582cdd259d15f26ae27241eab11ebe8b Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 22 May 2018 17:24:13 +0000 Subject: [PATCH 63/71] fix include. --- src/operator/control_flow.cc | 6 +++--- src/operator/subgraph_op_common.cc | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/operator/control_flow.cc b/src/operator/control_flow.cc index df28d5cb2cae..37d2b3168c6c 100644 --- a/src/operator/control_flow.cc +++ b/src/operator/control_flow.cc @@ -24,9 +24,9 @@ #include #include #include -#include "../operator_common.h" -#include "../elemwise_op_common.h" -#include "../../imperative/imperative_utils.h" +#include "./operator_common.h" +#include "./elemwise_op_common.h" +#include "../imperative/imperative_utils.h" #include "./subgraph_op_common.h" namespace mxnet { diff --git a/src/operator/subgraph_op_common.cc b/src/operator/subgraph_op_common.cc index ac2218b062fe..8344c24ab558 100644 --- a/src/operator/subgraph_op_common.cc +++ b/src/operator/subgraph_op_common.cc @@ -18,8 +18,8 @@ */ #include "./subgraph_op_common.h" -#include "../operator_common.h" -#include "../../imperative/imperative_utils.h" +#include "./operator_common.h" +#include "../imperative/imperative_utils.h" namespace mxnet { namespace op { From 471e6d2e58cb0b88b79335c5b42d4bbf5658c2ff Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 22 May 2018 17:24:31 +0000 Subject: [PATCH 64/71] add a test in gluon. --- tests/python/unittest/test_gluon_rnn.py | 62 ++++++++++++++++--------- 1 file changed, 41 insertions(+), 21 deletions(-) diff --git a/tests/python/unittest/test_gluon_rnn.py b/tests/python/unittest/test_gluon_rnn.py index cd885695412f..d4ac88900c5d 100644 --- a/tests/python/unittest/test_gluon_rnn.py +++ b/tests/python/unittest/test_gluon_rnn.py @@ -18,9 +18,10 @@ import mxnet as mx from mxnet import gluon import numpy as np +import copy from numpy.testing import assert_allclose import unittest -from mxnet.test_utils import almost_equal +from mxnet.test_utils import almost_equal, assert_almost_equal def test_rnn(): @@ -36,33 +37,52 @@ def test_rnn(): assert outs == [(10, 100), (10, 100), (10, 100)] -class RNNLayer(gluon.HybridBlock): - def __init__(self, prefix=None, params=None): - super(RNNLayer, self).__init__(prefix=prefix, params=params) - self.cell = gluon.contrib.rnn.RNNCell(100, prefix='rnn_') +class TestRNNLayer(gluon.HybridBlock): + def __init__(self, hidden_size, prefix=None, params=None): + super(TestRNNLayer, self).__init__(prefix=prefix, params=params) + self.cell = gluon.rnn.RNNCell(hidden_size, prefix='rnn_') - def hybrid_forward(self, F, inputs, states=None): - return self.cell.unroll(inputs, states) + def hybrid_forward(self, F, inputs, states): + states = [states] + out, states = F.contrib.foreach(self.cell, inputs, states) + return out def test_contrib_rnn(): - contrib_cell = gluon.contrib.rnn.RNNCell(100, prefix='rnn_') - inputs = mx.sym.Variable('rnn_data') - contrib_outputs, _ = contrib_cell.unroll(inputs) - assert sorted(contrib_cell.collect_params().keys()) == ['rnn_h2h_bias', 'rnn_h2h_weight', - 'rnn_i2h_bias', 'rnn_i2h_weight'] - - args, outs, auxs = contrib_outputs.infer_shape(rnn_data=(3, 10,50)) - assert outs == [(3, 10, 100)] - - rnn_data = mx.nd.normal(loc=0, scale=1, shape=(3, 10, 50)) - layer = RNNLayer() + batch_size = 10 + hidden_size = 100 + rnn_data = mx.nd.normal(loc=0, scale=1, shape=(5, batch_size, 50)) + states = mx.nd.normal(loc=0, scale=1, shape=(batch_size, hidden_size)) + layer = TestRNNLayer(hidden_size) layer.initialize(ctx=mx.cpu(0)) - res1 = layer(rnn_data) + res1 = layer(rnn_data, states) + params1 = layer.collect_params() + orig_params1 = copy.deepcopy(params1) + + trainer = gluon.Trainer(params1, 'sgd', {'learning_rate' : 0.03}) + with mx.autograd.record(): + res1 = layer(rnn_data, states) + res1.backward() + trainer.step(batch_size) - layer = RNNLayer() + layer = TestRNNLayer(hidden_size) layer.initialize(ctx=mx.cpu(0)) layer.hybridize() - res2 = layer(rnn_data) + res2 = layer(rnn_data, states) + params2 = layer.collect_params() + for key, val in orig_params1.items(): + params2[key].set_data(val.data()) + + trainer = gluon.Trainer(params2, 'sgd', {'learning_rate' : 0.03}) + with mx.autograd.record(): + res2 = layer(rnn_data, states) + assert_almost_equal(res1.asnumpy(), res2.asnumpy(), rtol=0.001, atol=0.0001) + res2.backward() + trainer.step(batch_size) + + for key, val in params1.items(): + weight1 = val.data() + weight2 = params2[key].data() + assert_almost_equal(weight1.asnumpy(), weight2.asnumpy(), rtol=0.001, atol=0.0001) def test_lstm(): From 580f294ecae2ed06332e43fc4d73b5fd538899be Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 22 May 2018 17:42:39 +0000 Subject: [PATCH 65/71] work for GPU. --- src/operator/control_flow.cc | 6 +++++- tests/python/unittest/test_operator.py | 8 ++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/operator/control_flow.cc b/src/operator/control_flow.cc index 37d2b3168c6c..8ae2e98c346a 100644 --- a/src/operator/control_flow.cc +++ b/src/operator/control_flow.cc @@ -517,6 +517,9 @@ NNVM_REGISTER_OP(_foreach) .set_attr("FInferShape", ForeachShape) .set_attr("FInferType", ForeachType) .set_attr("FStatefulComputeEx", ForeachComputeExCPU) +// Foreach operator works like an executor. Its code will always run on CPU. +// So the same code can be registered for both CPU and GPU. +.set_attr("FStatefulComputeEx", ForeachComputeExCPU) .set_attr("key_var_num_args", "num_args") .add_argument("fn", "Symbol", "Input graph.") .add_argument("data", "NDArray-or-Symbol[]", @@ -536,7 +539,8 @@ NNVM_REGISTER_OP(_backward_foreach) .set_attr_parser(ParamParser) .set_attr("TIsLayerOpBackward", true) .set_attr("TIsBackward", true) -.set_attr("FStatefulComputeEx", ForeachGradComputeExCPU); +.set_attr("FStatefulComputeEx", ForeachGradComputeExCPU) +.set_attr("FStatefulComputeEx", ForeachGradComputeExCPU); } // namespace op } // namespace mxnet diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 4fce81a7ae14..bb5fb50c25e0 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5719,7 +5719,7 @@ def verify_foreach(step, in_syms, state_syms, free_syms, name = name[1:] gin_order.append(int(name)) - e = out.bind(ctx=mx.cpu(), args=arg_dict, args_grad=arg_grad_dict) + e = out.bind(ctx=default_context(), args=arg_dict, args_grad=arg_grad_dict) e.forward(is_train=is_train) if (is_train): # backward @@ -5838,7 +5838,7 @@ def step(in1, states): state = mx.nd.arange(2) data_grad = mx.nd.empty(data.shape) state_grad = mx.nd.empty(state.shape) - e = out.bind(ctx=mx.cpu(), args={'v1':data, 'v2':state}, + e = out.bind(ctx=default_context(), args={'v1':data, 'v2':state}, args_grad={'v1':data_grad, 'v2':state_grad}) e.forward(is_train=True) out = mx.nd.zeros_like(data) @@ -5906,7 +5906,7 @@ def sym_group(out): h2h_barr_grad1 = mx.nd.empty(h2h_barr.shape) out = mx.sym.contrib.foreach(step, data, [init_h, init_c]) out = sym_group(out) - e1 = out.bind(ctx=mx.cpu(), + e1 = out.bind(ctx=default_context(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, 'i2h_weight': i2h_warr, 'h2h_weight': h2h_warr, 'i2h_bias': i2h_barr, 'h2h_bias': h2h_barr}, @@ -5937,7 +5937,7 @@ def sym_group(out): unroll_outs.append(mx.sym.expand_dims(h, axis=0)) unroll_outs = mx.sym.concat(*unroll_outs, dim=0) out = mx.sym.Group([unroll_outs, h, c]) - e2 = out.bind(ctx=mx.cpu(), + e2 = out.bind(ctx=default_context(), args={'data': data_arr, 'h': h_arr, 'c': c_arr, 'mylstm_i2h_weight': i2h_warr, 'mylstm_h2h_weight': h2h_warr, 'mylstm_i2h_bias': i2h_barr, 'mylstm_h2h_bias': h2h_barr}, From ae9340d70c68c7c6b7b0a111df736545ebfc5851 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 22 May 2018 17:58:10 +0000 Subject: [PATCH 66/71] small fix. --- src/executor/graph_executor.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/executor/graph_executor.cc b/src/executor/graph_executor.cc index cf1263b79057..ca06a12a5a0e 100644 --- a/src/executor/graph_executor.cc +++ b/src/executor/graph_executor.cc @@ -1557,7 +1557,7 @@ GraphExecutor::CachedSegOpr GraphExecutor::CreateCachedSegOpr(size_t topo_start, if (inode.source->is_variable()) continue; // We shouldn't add control flow operators to a segment. // We can't execute these operators in the engine. - if (op_node.exec->HasSubgraph()) continue; + if (op_node.exec->HasSubgraph()) return ret; if (op_node.exec->exec_type() != ExecType::kSync) { return ret; } From 977f5624ad0b0dedb9dcb8629f975afc56bb1e1a Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 22 May 2018 18:02:57 +0000 Subject: [PATCH 67/71] remove subgraph_name --- python/mxnet/symbol/contrib.py | 19 +++++++++---------- src/c_api/c_api_symbolic.cc | 5 ++--- 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index 2d067555e8dc..c1a04a105d74 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -187,16 +187,15 @@ def check_data(inputs, in_type, msg): # the python function, we need to prune the computation graph constructed from # the function. One way of doing it is to mark the nodes in the computation graph # with AttrScope and prune the nodes without the special attribute. - with AttrScope(subgraph_name=name): - if isinstance(data, list): - in_eles = [symbol.var(sym.name) for sym in data] - else: - in_eles = symbol.var(data.name) - if isinstance(init_states, list): - states = [symbol.var(s.name) for s in init_states] - else: - states = symbol.var(init_states.name) - sym_out, sym_states = body(in_eles, states) + if isinstance(data, list): + in_eles = [symbol.var(sym.name) for sym in data] + else: + in_eles = symbol.var(data.name) + if isinstance(init_states, list): + states = [symbol.var(s.name) for s in init_states] + else: + states = symbol.var(init_states.name) + sym_out, sym_states = body(in_eles, states) check_data(sym_out, symbol.Symbol, "the output should be an NDArray or a list of NDArrays") check_data(sym_states, symbol.Symbol, diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index 34819e83d506..be008e92e86f 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -38,11 +38,10 @@ void RegisterLegacyOpProp(); void RegisterLegacyNDFunc(); } const std::vector kHiddenKeys = { - "ctx_group", "lr_mult", "wd_mult", "force_mirroring", "mirror_stage", "subgraph_name" + "ctx_group", "lr_mult", "wd_mult", "force_mirroring", "mirror_stage" }; const std::vector kReplacedHiddenKeys = { - "__ctx_group__", "__lr_mult__", "__wd_mult__", "__force_mirroring__", "__mirror_stage__", - "subgraph_name" + "__ctx_group__", "__lr_mult__", "__wd_mult__", "__force_mirroring__", "__mirror_stage__" }; const char *kNamespaceSeparator = "$"; From fab06e526d798a9ce583cad7152ad1100bfd04f5 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 22 May 2018 18:32:41 +0000 Subject: [PATCH 68/71] create loop state for reuse in the future. --- src/operator/control_flow.cc | 131 +---------------------------- src/operator/subgraph_op_common.cc | 101 ++++++++++++++++++++++ src/operator/subgraph_op_common.h | 37 ++++++++ 3 files changed, 140 insertions(+), 129 deletions(-) diff --git a/src/operator/control_flow.cc b/src/operator/control_flow.cc index 37d2b3168c6c..0210bd45c80a 100644 --- a/src/operator/control_flow.cc +++ b/src/operator/control_flow.cc @@ -57,142 +57,15 @@ struct ForeachParam : public dmlc::Parameter { DMLC_REGISTER_PARAMETER(ForeachParam); -class ForeachState { - // These are output arrays from all iterations. - // They also contain the Op state for each CachedOp. - std::vector > all_outputs; - std::vector > all_inputs; - std::vector > all_gradients; - std::vector iter_ops; - +class ForeachState: public LoopState { public: - Symbol subgraph_sym; - nnvm::Graph subgraph; ForeachParam params; - ForeachState(const Symbol &g, const ForeachParam ¶ms) { - this->subgraph_sym = g; - this->subgraph.outputs = g.outputs; + ForeachState(const Symbol &g, const ForeachParam ¶ms) : LoopState(g) { this->params = params; } - - void Forward(std::vector cinputs, - const std::vector& req, - std::vector coutputs, - bool is_recording); - void Backward(int iter_no, - std::vector ograds, - const std::vector &req, - std::vector igrads); - void Cleanup() { - all_outputs.clear(); - all_inputs.clear(); - all_gradients.clear(); - iter_ops.clear(); - } }; -void ForeachState::Forward(std::vector cinputs, - const std::vector& req, - std::vector coutputs, - bool is_recording) { - using namespace nnvm; - using namespace imperative; - - bool orig_is_record; - if (is_recording) - orig_is_record = Imperative::Get()->set_is_recording(true); - else - orig_is_record = Imperative::Get()->is_recording(); - - std::vector inputs(cinputs.size()); - std::vector outputs(coutputs.size()); - for (size_t i = 0; i < inputs.size(); i++) - inputs[i] = &cinputs[i]; - for (size_t i = 0; i < outputs.size(); i++) - outputs[i] = &coutputs[i]; - - if (is_recording) { - all_inputs.push_back(cinputs); - std::vector gradients(cinputs.size()); - std::vector input_ptrs(cinputs.size()); - std::vector gradient_ptrs(cinputs.size()); - std::vector grad_reqs(cinputs.size()); - for (size_t i = 0; i < gradients.size(); i++) { - gradients[i] = NDArray(cinputs[i].shape(), cinputs[i].ctx(), - true, cinputs[i].dtype()); - input_ptrs[i] = &cinputs[i]; - gradient_ptrs[i] = &gradients[i]; - grad_reqs[i] = kWriteTo; - } - Imperative::Get()->MarkVariables(input_ptrs, grad_reqs, gradient_ptrs);; - } - - std::vector > kwargs; - kwargs.push_back(std::pair("inline_limit", "0")); - // Get input names. - const auto& idx = subgraph.indexed_graph(); - std::vector arg_names(idx.input_nodes().size()); - for (size_t i = 0; i < idx.input_nodes().size(); ++i) - arg_names[i] = idx[idx.input_nodes()[i]].source->attrs.name; - // We don't have parameters for the cached op. - std::unordered_map > params; - CachedOpPtr op = std::make_shared(subgraph_sym, kwargs, - arg_names, params); - // TODO(zhengda) we need to avoid shape inference and memory plan whenever the op is - // called. Currently, CachedOp allocates memory each time Forward is called. - // I need to fix this once the PR for static memory allocation in CachedOp is - // merged. https://github.com/apache/incubator-mxnet/pull/10817 - op->Forward(nullptr, inputs, outputs); - - if (is_recording) { - all_outputs.push_back(coutputs); - iter_ops.push_back(op); - } - - Imperative::Get()->set_is_recording(orig_is_record); -} - -void ForeachState::Backward(int iter_no, - std::vector ograds, - const std::vector &req, - std::vector igrads) { - using namespace nnvm; - using namespace imperative; - - CHECK_GT(iter_ops.size(), iter_no) - << "We didn't record the computation for iteration " << iter_no; - auto op = iter_ops[iter_no]; - std::vector inputs; - std::vector outputs; - inputs.reserve(op->num_backward_inputs()); - outputs.reserve(op->num_inputs()); - for (size_t i = 0; i < ograds.size(); i++) - inputs.push_back(&ograds[i]); - - const std::vector &save_inputs = op->save_inputs(); - const std::vector &save_outputs = op->save_outputs(); - CHECK_EQ(save_inputs.size(), all_inputs[iter_no].size()); - CHECK_EQ(op->num_outputs(), all_outputs[iter_no].size()); - for (size_t i = 0; i < all_inputs[iter_no].size(); i++) { - if (save_inputs[i]) - inputs.push_back(&all_inputs[iter_no][i]); - } - for (size_t i = 0; i < all_outputs[iter_no].size(); i++) { - if (save_outputs[i]) - inputs.push_back(&all_outputs[iter_no][i]); - } - CHECK_EQ(inputs.size(), op->num_backward_inputs()); - for (size_t i = 0; i < igrads.size(); i++) - outputs.push_back(&igrads[i]); - CHECK_EQ(outputs.size(), op->num_inputs()); - - CHECK(!Imperative::AGInfo::IsNone(all_outputs[iter_no][0])); - const nnvm::NodeEntry &node_entry = all_outputs[iter_no][0].entry(); - OpStatePtr state = Imperative::AGInfo::Get(node_entry.node).state; - op->Backward(false, state, inputs, req, outputs); -} - static void ForeachComputeExCPU(const OpStatePtr& state_ptr, const OpContext& ctx, const std::vector& inputs, diff --git a/src/operator/subgraph_op_common.cc b/src/operator/subgraph_op_common.cc index 8344c24ab558..fa22898c13d4 100644 --- a/src/operator/subgraph_op_common.cc +++ b/src/operator/subgraph_op_common.cc @@ -151,5 +151,106 @@ bool InferSubgraphBackwardStorage(const nnvm::Symbol &subgraph, return true; } +void LoopState::Forward(std::vector cinputs, + const std::vector& req, + std::vector coutputs, + bool is_recording) { + using namespace nnvm; + using namespace imperative; + + bool orig_is_record; + if (is_recording) + orig_is_record = Imperative::Get()->set_is_recording(true); + else + orig_is_record = Imperative::Get()->is_recording(); + + std::vector inputs(cinputs.size()); + std::vector outputs(coutputs.size()); + for (size_t i = 0; i < inputs.size(); i++) + inputs[i] = &cinputs[i]; + for (size_t i = 0; i < outputs.size(); i++) + outputs[i] = &coutputs[i]; + + if (is_recording) { + all_inputs.push_back(cinputs); + std::vector gradients(cinputs.size()); + std::vector input_ptrs(cinputs.size()); + std::vector gradient_ptrs(cinputs.size()); + std::vector grad_reqs(cinputs.size()); + for (size_t i = 0; i < gradients.size(); i++) { + gradients[i] = NDArray(cinputs[i].shape(), cinputs[i].ctx(), + true, cinputs[i].dtype()); + input_ptrs[i] = &cinputs[i]; + gradient_ptrs[i] = &gradients[i]; + grad_reqs[i] = kWriteTo; + } + Imperative::Get()->MarkVariables(input_ptrs, grad_reqs, gradient_ptrs);; + } + + std::vector > kwargs; + kwargs.push_back(std::pair("inline_limit", "0")); + // Get input names. + const auto& idx = subgraph.indexed_graph(); + std::vector arg_names(idx.input_nodes().size()); + for (size_t i = 0; i < idx.input_nodes().size(); ++i) + arg_names[i] = idx[idx.input_nodes()[i]].source->attrs.name; + // We don't have parameters for the cached op. + std::unordered_map > params; + CachedOpPtr op = std::make_shared(subgraph_sym, kwargs, + arg_names, params); + // TODO(zhengda) we need to avoid shape inference and memory plan whenever the op is + // called. Currently, CachedOp allocates memory each time Forward is called. + // I need to fix this once the PR for static memory allocation in CachedOp is + // merged. https://github.com/apache/incubator-mxnet/pull/10817 + op->Forward(nullptr, inputs, outputs); + + if (is_recording) { + all_outputs.push_back(coutputs); + iter_ops.push_back(op); + } + + Imperative::Get()->set_is_recording(orig_is_record); +} + +void LoopState::Backward(int iter_no, + std::vector ograds, + const std::vector &req, + std::vector igrads) { + using namespace nnvm; + using namespace imperative; + + CHECK_GT(iter_ops.size(), iter_no) + << "We didn't record the computation for iteration " << iter_no; + auto op = iter_ops[iter_no]; + std::vector inputs; + std::vector outputs; + inputs.reserve(op->num_backward_inputs()); + outputs.reserve(op->num_inputs()); + for (size_t i = 0; i < ograds.size(); i++) + inputs.push_back(&ograds[i]); + + const std::vector &save_inputs = op->save_inputs(); + const std::vector &save_outputs = op->save_outputs(); + CHECK_EQ(save_inputs.size(), all_inputs[iter_no].size()); + CHECK_EQ(op->num_outputs(), all_outputs[iter_no].size()); + for (size_t i = 0; i < all_inputs[iter_no].size(); i++) { + if (save_inputs[i]) + inputs.push_back(&all_inputs[iter_no][i]); + } + for (size_t i = 0; i < all_outputs[iter_no].size(); i++) { + if (save_outputs[i]) + inputs.push_back(&all_outputs[iter_no][i]); + } + CHECK_EQ(inputs.size(), op->num_backward_inputs()); + for (size_t i = 0; i < igrads.size(); i++) + outputs.push_back(&igrads[i]); + CHECK_EQ(outputs.size(), op->num_inputs()); + + CHECK(!Imperative::AGInfo::IsNone(all_outputs[iter_no][0])); + const nnvm::NodeEntry &node_entry = all_outputs[iter_no][0].entry(); + OpStatePtr state = Imperative::AGInfo::Get(node_entry.node).state; + op->Backward(false, state, inputs, req, outputs); +} + } // namespace op } // namespace mxnet diff --git a/src/operator/subgraph_op_common.h b/src/operator/subgraph_op_common.h index 25cbd60f5b63..74e7cb2d1ccd 100644 --- a/src/operator/subgraph_op_common.h +++ b/src/operator/subgraph_op_common.h @@ -24,6 +24,7 @@ #include #include #include +#include "../imperative/imperative_utils.h" namespace mxnet { namespace op { @@ -55,6 +56,42 @@ bool InferSubgraphBackwardStorage(const nnvm::Symbol &subgraph, std::vector *in_attrs, std::vector *out_attrs); +/* + * This contains the states for running a loop and provides methods + * of running the subgraph computation for an iteration. + */ +class LoopState { + // These are output arrays from all iterations. + // They also contain the Op state for each CachedOp. + std::vector > all_outputs; + std::vector > all_inputs; + std::vector > all_gradients; + std::vector iter_ops; + Symbol subgraph_sym; + nnvm::Graph subgraph; + + public: + LoopState(const Symbol &g) { + this->subgraph_sym = g; + this->subgraph.outputs = g.outputs; + } + + void Forward(std::vector cinputs, + const std::vector& req, + std::vector coutputs, + bool is_recording); + void Backward(int iter_no, + std::vector ograds, + const std::vector &req, + std::vector igrads); + void Cleanup() { + all_outputs.clear(); + all_inputs.clear(); + all_gradients.clear(); + iter_ops.clear(); + } +}; + } // namespace op } // namespace mxnet From 08fbd04f1a7bb187b54b2073b8a115d2ffbf9364 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Tue, 22 May 2018 22:07:05 +0000 Subject: [PATCH 69/71] move code. --- src/c_api/c_api_symbolic.cc | 21 ++++----------- src/nnvm/graph_editor.cc | 51 +++++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 16 deletions(-) create mode 100644 src/nnvm/graph_editor.cc diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index be008e92e86f..b3b9cbc61ee6 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -344,26 +344,15 @@ int MXSymbolGetAtomicSymbolName(AtomicSymbolCreator creator, API_END(); } +namespace mxnet { +extern std::vector GetInputSymbols(const nnvm::Symbol &sym); +} + int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **input_arr, int *input_size) { API_BEGIN(); nnvm::Symbol *s = static_cast(sym); - nnvm::Graph g; - g.outputs = s->outputs; - std::vector input_syms; - const nnvm::IndexedGraph& idx = g.indexed_graph(); size_t max_input_size = *input_size; - // Go through all nodes and return the ones representing variables. - for (size_t i = 0; i < idx.num_nodes(); i++) { - const nnvm::Node &n = *idx[i].source; - for (const nnvm::NodeEntry &e : n.inputs) { - auto p = e.node; - if (p->is_variable()) { - nnvm::Symbol *s = new nnvm::Symbol(); - s->outputs.push_back(e); - input_syms.push_back(s); - } - } - } + std::vector input_syms = mxnet::GetInputSymbols(*s); CHECK(input_syms.size() <= max_input_size); *input_size = input_syms.size(); memcpy(input_arr, input_syms.data(), sizeof(*input_arr) * input_syms.size()); diff --git a/src/nnvm/graph_editor.cc b/src/nnvm/graph_editor.cc new file mode 100644 index 000000000000..7200f5b124e4 --- /dev/null +++ b/src/nnvm/graph_editor.cc @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file graph_editor.cc + * The functions in this file edit an NNVM graph. Potentially, + * these functions should be moved to NNVM in the future. + */ + +#include +#include + +namespace mxnet { + +std::vector GetInputSymbols(const nnvm::Symbol &sym) { + nnvm::Graph g; + std::vector input_syms; + g.outputs = sym.outputs; + const nnvm::IndexedGraph& idx = g.indexed_graph(); + // Go through all nodes and return the ones representing variables. + for (size_t i = 0; i < idx.num_nodes(); i++) { + const nnvm::Node &n = *idx[i].source; + for (const nnvm::NodeEntry &e : n.inputs) { + auto p = e.node; + if (p->is_variable()) { + nnvm::Symbol *s = new nnvm::Symbol(); + s->outputs.push_back(e); + input_syms.push_back(s); + } + } + } + return input_syms; +} + +} From 4550f8521098831dec364064f925db8e89cdfa5b Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Wed, 23 May 2018 23:32:00 +0000 Subject: [PATCH 70/71] Revert "remove subgraph_name" This reverts commit 977f5624ad0b0dedb9dcb8629f975afc56bb1e1a. --- python/mxnet/symbol/contrib.py | 19 ++++++++++--------- src/c_api/c_api_symbolic.cc | 5 +++-- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index c1a04a105d74..2d067555e8dc 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -187,15 +187,16 @@ def check_data(inputs, in_type, msg): # the python function, we need to prune the computation graph constructed from # the function. One way of doing it is to mark the nodes in the computation graph # with AttrScope and prune the nodes without the special attribute. - if isinstance(data, list): - in_eles = [symbol.var(sym.name) for sym in data] - else: - in_eles = symbol.var(data.name) - if isinstance(init_states, list): - states = [symbol.var(s.name) for s in init_states] - else: - states = symbol.var(init_states.name) - sym_out, sym_states = body(in_eles, states) + with AttrScope(subgraph_name=name): + if isinstance(data, list): + in_eles = [symbol.var(sym.name) for sym in data] + else: + in_eles = symbol.var(data.name) + if isinstance(init_states, list): + states = [symbol.var(s.name) for s in init_states] + else: + states = symbol.var(init_states.name) + sym_out, sym_states = body(in_eles, states) check_data(sym_out, symbol.Symbol, "the output should be an NDArray or a list of NDArrays") check_data(sym_states, symbol.Symbol, diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index b3b9cbc61ee6..1a5c00630580 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -38,10 +38,11 @@ void RegisterLegacyOpProp(); void RegisterLegacyNDFunc(); } const std::vector kHiddenKeys = { - "ctx_group", "lr_mult", "wd_mult", "force_mirroring", "mirror_stage" + "ctx_group", "lr_mult", "wd_mult", "force_mirroring", "mirror_stage", "subgraph_name" }; const std::vector kReplacedHiddenKeys = { - "__ctx_group__", "__lr_mult__", "__wd_mult__", "__force_mirroring__", "__mirror_stage__" + "__ctx_group__", "__lr_mult__", "__wd_mult__", "__force_mirroring__", "__mirror_stage__", + "subgraph_name" }; const char *kNamespaceSeparator = "$"; From d3c4f6f85cdb85c300269de6e372547b93b525d7 Mon Sep 17 00:00:00 2001 From: Da Zheng Date: Fri, 25 May 2018 02:38:57 +0000 Subject: [PATCH 71/71] cut graph. --- include/mxnet/c_api.h | 16 +++++- python/mxnet/symbol/contrib.py | 67 +++++++++++++++++--------- src/c_api/c_api_symbolic.cc | 54 +++++++++++++++++++++ src/nnvm/graph_editor.cc | 34 +++++++++++++ tests/python/unittest/test_operator.py | 22 +++++++-- 5 files changed, 165 insertions(+), 28 deletions(-) diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h index a8e87271e166..79c92edfa0a1 100644 --- a/include/mxnet/c_api.h +++ b/include/mxnet/c_api.h @@ -1060,12 +1060,24 @@ MXNET_DLL int MXSymbolGetAtomicSymbolName(AtomicSymbolCreator creator, /*! * \brief Get the input symbols of the graph. * \param sym The graph. - * \param outs The input symbols of the graph. - * \param out_size the number of input symbols returned. + * \param inputs The input symbols of the graph. + * \param input_size the number of input symbols returned. */ MXNET_DLL int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **inputs, int *input_size); +/*! + * \brief Cut a subgraph whose nodes are marked with a subgraph attribute. + * The input graph will be modified. A variable node will be created for each + * edge that connects to nodes outside the subgraph. The outside nodes that + * connect to the subgraph will be returned. + * \param sym The graph. + * \param inputs The nodes that connect to the subgraph. + * \param input_size The number of such nodes. + */ +MXNET_DLL int MXSymbolCutSubgraph(SymbolHandle sym, SymbolHandle **inputs, + int *input_size); + /*! * \brief Get the detailed information about atomic symbol. * \param creator the AtomicSymbolCreator. diff --git a/python/mxnet/symbol/contrib.py b/python/mxnet/symbol/contrib.py index 2d067555e8dc..a1a1d23bbbe5 100644 --- a/python/mxnet/symbol/contrib.py +++ b/python/mxnet/symbol/contrib.py @@ -20,6 +20,7 @@ """Contrib Symbol API of MXNet.""" import math import ctypes +import re from .random import uniform from .symbol import Symbol @@ -110,6 +111,17 @@ def _get_graph_inputs(subg): syms.append(s) return syms +def _cut_subgraph(subg): + num_handles = ctypes.c_int(1000) + handles = c_array(SymbolHandle, [SymbolHandle(0) for i in range(1000)]) + check_call(_LIB.MXSymbolCutSubgraph(subg.handle, handles, ctypes.byref(num_handles))) + + syms = [] + for i in range(num_handles.value): + s = Symbol(handles[i]) + syms.append(s) + return syms + def foreach(body, data, init_states, name="foreach"): """Run a for loop with user-defined computation over Symbols on dimension 0. @@ -198,28 +210,31 @@ def check_data(inputs, in_type, msg): states = symbol.var(init_states.name) sym_out, sym_states = body(in_eles, states) - check_data(sym_out, symbol.Symbol, "the output should be an NDArray or a list of NDArrays") - check_data(sym_states, symbol.Symbol, - "the output states should be an NDArray or a list of NDArrays") - if isinstance(sym_states, list): - assert isinstance(init_states, list) and len(sym_states) == len(init_states), \ - "the number of output states (%d) should be the same as input states (%d)" \ - % (len(sym_states), len(init_states)) + check_data(sym_out, symbol.Symbol, + "the output should be an NDArray or a list of NDArrays") + check_data(sym_states, symbol.Symbol, + "the output states should be an NDArray or a list of NDArrays") + if isinstance(sym_states, list): + assert isinstance(init_states, list) and len(sym_states) == len(init_states), \ + "the number of output states (%d) should be the same as input states (%d)" \ + % (len(sym_states), len(init_states)) + + if isinstance(sym_out, list): + flat_out = sym_out + else: + flat_out = [sym_out] + num_out_data = len(flat_out) + if isinstance(sym_states, list): + for s in sym_states: + # There is a problem if the outputs are the same as the inputs + # or the first output. By calling identity, we can make sure that + # all symbols will refer to different NDArrays. + flat_out.append(symbol.op.identity(s)) + else: + flat_out.append(symbol.op.identity(sym_states)) + g = symbol.Group(flat_out) - if isinstance(sym_out, list): - flat_out = sym_out - else: - flat_out = [sym_out] - num_out_data = len(flat_out) - if isinstance(sym_states, list): - for s in sym_states: - # There is a problem if the outputs are the same as the inputs - # or the first output. By calling identity, we can make sure that - # all symbols will refer to different NDArrays. - flat_out.append(symbol.op.identity(s)) - else: - flat_out.append(symbol.op.identity(sym_states)) - g = symbol.Group(flat_out) + cut_syms = _cut_subgraph(g) input_syms = _get_graph_inputs(g) # Here we need to find out how the input symbols are ordered as well as @@ -230,12 +245,13 @@ def check_data(inputs, in_type, msg): gin_names = input_syms.keys() # This array contains the symbols for the inputs of foreach. # They are ordered according to the inputs of the subgraph. - ordered_ins = [] states_map = {sym.name:sym for sym in init_states} state_names = states_map.keys() data_syms = _as_list(data) data_map = {sym.name:sym for sym in data_syms} data_names = data_map.keys() + + ordered_ins = [] in_state_locs = [] in_data_locs = [] for in_name in g.list_inputs(): @@ -248,7 +264,12 @@ def check_data(inputs, in_type, msg): ordered_ins.append(data_map[in_name]) in_data_locs.append(len(ordered_ins) - 1) else: - ordered_ins.append(input_syms[in_name]) + # The remaining inputs are the ones cut from the original graph. + # The names of these variable nodes contain the index in cut_syms. + m = re.search(r'\d+$', in_name) + idx = int(m.group()) if m else None + assert idx < len(cut_syms) + ordered_ins.append(cut_syms[idx]) num_outputs = len(flat_out) num_states = len(state_names) diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc index 1a5c00630580..030ab432228b 100644 --- a/src/c_api/c_api_symbolic.cc +++ b/src/c_api/c_api_symbolic.cc @@ -346,7 +346,13 @@ int MXSymbolGetAtomicSymbolName(AtomicSymbolCreator creator, } namespace mxnet { + extern std::vector GetInputSymbols(const nnvm::Symbol &sym); +extern bool CutGraph(const std::vector &input_entries, + const std::string &in_name_prefix, bool skip_var, + std::vector *orig_entries, + std::vector *new_var_names); + } int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **input_arr, int *input_size) { @@ -360,6 +366,54 @@ int MXSymbolGetInputSymbols(SymbolHandle sym, SymbolHandle **input_arr, int *inp API_END_HANDLE_ERROR(); } +int MXSymbolCutSubgraph(SymbolHandle sym, SymbolHandle **input_symbols, + int *input_size) { + // Given a graph, we want to fetch the nodes that have been marked as part of + // a subgraph. + API_BEGIN(); + nnvm::Symbol *s = static_cast(sym); + size_t max_input_size = *input_size; + std::string subg_attr = "__subgraph_name__"; + auto out_node = s->outputs[0].node; + auto it = out_node->attrs.dict.find(subg_attr); + if (it != out_node->attrs.dict.end()) { + std::string subg_name = it->second; + std::vector input_entries; + DFSVisit(s->outputs, [subg_attr, subg_name, &input_entries] + (nnvm::NodePtr n) { + // If the node itself isn't in the subgraph, we ignore it. + auto it = n->attrs.dict.find(subg_attr); + if (it == n->attrs.dict.end() || it->second != subg_name) + return; + + // We search for nodes whose node entries aren't in the subgraph. + for (size_t j = 0; j < n->inputs.size(); j++) { + auto in_node = n->inputs[j].node; + auto it = in_node->attrs.dict.find(subg_attr); + if (it == in_node->attrs.dict.end() || it->second != subg_name) + input_entries.push_back(&n->inputs[j]); + } + }); + + std::vector orig_entries; + std::vector new_var_names; + CutGraph(input_entries, subg_name + "_var", false, &orig_entries, &new_var_names); + + std::vector input_syms(orig_entries.size()); + for (size_t i = 0; i < input_syms.size(); i++) { + input_syms[i] = new nnvm::Symbol(); + input_syms[i]->outputs.push_back(orig_entries[i]); + } + CHECK(input_syms.size() <= max_input_size); + *input_size = input_syms.size(); + memcpy(input_symbols, input_syms.data(), sizeof(*input_symbols) * input_syms.size()); + } else { + *input_size = 0; + } + + API_END_HANDLE_ERROR(); +} + int MXSymbolCreateFromFile(const char *fname, SymbolHandle *out) { nnvm::Symbol *s = new nnvm::Symbol(); API_BEGIN(); diff --git a/src/nnvm/graph_editor.cc b/src/nnvm/graph_editor.cc index 7200f5b124e4..98c99e2425df 100644 --- a/src/nnvm/graph_editor.cc +++ b/src/nnvm/graph_editor.cc @@ -25,9 +25,18 @@ #include #include +#include + +namespace nnvm { +NodePtr CreateVariableNode(const std::string& name); +} namespace mxnet { +/* + * Given a computation graph, this function finds the input nodes of the graph + * and create symbols for the input nodes. It returns the input symbols. + */ std::vector GetInputSymbols(const nnvm::Symbol &sym) { nnvm::Graph g; std::vector input_syms; @@ -48,4 +57,29 @@ std::vector GetInputSymbols(const nnvm::Symbol &sym) { return input_syms; } +/* + * Given a computation graph and a set of input node entries, this function cuts + * the node entries and creates new variable nodes as the input nodes of the + * subgraph. It returns the nodes that connect to the subgraph directly and + * the names of the new variable nodes. + */ +bool CutGraph(const std::vector &input_entries, + const std::string &in_name_prefix, bool skip_var, + std::vector *orig_entries, + std::vector *new_var_names) { + orig_entries->reserve(input_entries.size()); + for (size_t i = 0; i < input_entries.size(); i++) { + nnvm::NodeEntry *e = input_entries[i]; + // If the node is a variable itself, we may want to skip the node. + if (e->node->is_variable() && skip_var) + continue; + + orig_entries->push_back(*e); + new_var_names->push_back(in_name_prefix + std::to_string(i)); + nnvm::NodePtr n = nnvm::CreateVariableNode(new_var_names->back()); + *e = nnvm::NodeEntry{n, 0, 0}; + } + return true; +} + } diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index bb5fb50c25e0..2b2a66725bb1 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -5683,7 +5683,8 @@ def step3(in1, states, free): return ([out, out * 2], [out * 2, out * 3]) def verify_foreach(step, in_syms, state_syms, free_syms, - in_arrs, init_states, frees, out_grads, is_train=True): + in_arrs, init_states, frees, out_grads, is_train=True, + free_vars_func=None): step_sym = lambda in_syms, state_syms : step(in_syms, state_syms, free_syms) res, states = mx.sym.contrib.foreach(step_sym, in_syms, state_syms) out = _as_list(res) @@ -5737,8 +5738,9 @@ def verify_foreach(step, in_syms, state_syms, free_syms, arr.attach_grad() for arr in frees: arr.attach_grad() - step_imp = lambda in_arrs, state_arrs : step(in_arrs, state_arrs, frees) with mx.autograd.record(): + frees_imp = frees if free_vars_func is None else free_vars_func(frees) + step_imp = lambda in_arrs, state_arrs : step(in_arrs, state_arrs, frees_imp) states = [mx.nd.expand_dims(s, 0) for s in init_states] res, states = mx.nd.contrib.foreach(step_imp, in_arrs, init_states) @@ -5777,7 +5779,21 @@ def verify_foreach(step, in_syms, state_syms, free_syms, # * multiple inputs and multiple outputs. # * inference. - states = [mx.nd.random.uniform(shape=(2))] + #states = [mx.nd.random.uniform(shape=(2))] + + #frees1 = [mx.nd.random.uniform(shape=(2)), mx.nd.random.uniform(shape=(2))] + #arrs = mx.nd.random.uniform(shape=(3, 2)) + states = [mx.nd.arange(2)] + + frees1 = [mx.nd.arange(2), mx.nd.arange(2) + 1] + arrs = mx.nd.arange(6).reshape(shape=(3, 2)) + out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], + [mx.nd.random.uniform(-10, 10, states[0].shape)]] + verify_foreach(step1, v3, [v4], [v5 + v6], arrs, states, frees1, out_grads, True, + lambda frees : [frees[0] + frees[1]]) + verify_foreach(step1, v3, [v4], [v5 + v6], arrs, states, frees1, out_grads, False, + lambda frees : [frees[0] + frees[1]]) + frees = [mx.nd.random.uniform(shape=(2))] arrs = mx.nd.random.uniform(shape=(2, 2)) out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)],