You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by ma...@apache.org on 2018/06/16 01:27:39 UTC
[incubator-mxnet] branch revert-11313-static created (now 24f14f7)
This is an automated email from the ASF dual-hosted git repository.
marcoabreu pushed a change to branch revert-11313-static
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git.
at 24f14f7 Revert "Static alloc for hybridblock (#11313)"
This branch includes the following new commits:
new 24f14f7 Revert "Static alloc for hybridblock (#11313)"
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
--
To stop receiving notification emails like this one, please contact
marcoabreu@apache.org.
[incubator-mxnet] 01/01: Revert "Static alloc for hybridblock
(#11313)"
Posted by ma...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
marcoabreu pushed a commit to branch revert-11313-static
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
commit 24f14f7fdbef427aab25738608267c57de645987
Author: Marco de Abreu <ma...@users.noreply.github.com>
AuthorDate: Fri Jun 15 18:27:32 2018 -0700
Revert "Static alloc for hybridblock (#11313)"
This reverts commit 5431e12f11fc5446f6ec2a25098b4e3b67ee7eb3.
---
include/mxnet/c_api.h | 5 +
include/mxnet/imperative.h | 89 ++++
include/mxnet/ndarray.h | 8 -
include/mxnet/op_attr_types.h | 33 +-
python/mxnet/_ctypes/ndarray.py | 16 +-
python/mxnet/gluon/block.py | 74 ++--
src/c_api/c_api_ndarray.cc | 26 +-
src/engine/threaded_engine.cc | 3 +-
src/executor/attach_op_execs_pass.cc | 165 ++++---
src/executor/attach_op_resource_pass.cc | 16 +-
src/executor/exec_pass.h | 28 +-
src/executor/graph_executor.cc | 2 +-
src/imperative/cached_op.cc | 750 ++++++--------------------------
src/imperative/cached_op.h | 174 --------
src/imperative/imperative.cc | 90 +++-
src/imperative/imperative_utils.cc | 120 -----
src/imperative/imperative_utils.h | 256 ++---------
tests/python/unittest/test_gluon.py | 67 +--
18 files changed, 523 insertions(+), 1399 deletions(-)
diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h
index 4dd858a..55c26bc 100644
--- a/include/mxnet/c_api.h
+++ b/include/mxnet/c_api.h
@@ -987,6 +987,11 @@ MXNET_DLL int MXCreateCachedOpEx(SymbolHandle handle,
int num_flags,
const char** keys,
const char** vals,
+ int num_inputs,
+ const char** input_names,
+ int num_params,
+ const char** param_names,
+ NDArrayHandle* params,
CachedOpHandle *out);
/*!
* \brief free cached operator
diff --git a/include/mxnet/imperative.h b/include/mxnet/imperative.h
index 7ea60df..758ce85 100644
--- a/include/mxnet/imperative.h
+++ b/include/mxnet/imperative.h
@@ -35,6 +35,23 @@
#include "./ndarray.h"
namespace mxnet {
+/*! \brief CachedOp Parameters */
+struct CachedOpConfig : public dmlc::Parameter<CachedOpConfig> {
+ uint32_t inline_limit;
+ uint32_t forward_bulk_size;
+ uint32_t backward_bulk_size;
+ DMLC_DECLARE_PARAMETER(CachedOpConfig) {
+ DMLC_DECLARE_FIELD(inline_limit)
+ .set_default(2)
+ .describe("Maximum number of operators that can be inlined.");
+ DMLC_DECLARE_FIELD(forward_bulk_size)
+ .set_default(dmlc::GetEnv("MXNET_EXEC_BULK_EXEC_MAX_NODE_TRAIN", 15))
+ .describe("Segment size of bulk execution during forward pass.");
+ DMLC_DECLARE_FIELD(backward_bulk_size)
+ .set_default(dmlc::GetEnv("MXNET_EXEC_BULK_EXEC_MAX_NODE_TRAIN", 15))
+ .describe("Segment size of bulk execution during backward pass.");
+ }
+};
/*! \brief runtime functions for NDArray */
class Imperative {
public:
@@ -77,6 +94,67 @@ class Imperative {
&& info.out_grads.size() == 1;
}
};
+ class CachedOp {
+ public:
+ CachedOp(
+ const nnvm::Symbol& sym,
+ const std::vector<std::pair<std::string, std::string> >& flags,
+ const std::vector<std::string> arg_names,
+ const std::unordered_map<std::string, std::vector<NDArray> >& params);
+ uint32_t num_inputs() {
+ return fwd_graph_.indexed_graph().input_nodes().size();
+ }
+ uint32_t num_outputs() {
+ return fwd_graph_.outputs.size();
+ }
+ uint32_t num_backward_inputs() {
+ return bwd_ograd_dep_.size() + bwd_in_dep_.size() + bwd_out_dep_.size();
+ }
+ std::vector<bool>& save_inputs() {
+ return save_inputs_;
+ }
+ std::vector<bool>& save_outputs() {
+ return save_outputs_;
+ }
+ const std::unordered_set<uint32_t>& mutable_input_nodes() {
+ return fwd_graph_.indexed_graph().mutable_input_nodes();
+ }
+ nnvm::Graph GetForwardGraph(const bool recording,
+ const std::vector<NDArray*>& inputs);
+ nnvm::Graph GetBackwardGraph(const OpStatePtr& state,
+ const std::vector<OpReqType>& reqs,
+ const std::vector<NDArray*>& inputs);
+ std::vector<nnvm::NodeEntry> Gradient(const nnvm::NodePtr& node,
+ const std::vector<nnvm::NodeEntry>& ograds);
+ void Forward(const std::shared_ptr<CachedOp>& op_ptr,
+ const std::vector<NDArray*>& args,
+ const std::vector<NDArray*>& outputs);
+ void Backward(const bool retain_graph,
+ const OpStatePtr& state,
+ const std::vector<NDArray*>& inputs,
+ const std::vector<OpReqType>& reqs,
+ const std::vector<NDArray*>& outputs);
+
+ private:
+ struct CachedOpState {
+ std::vector<NDArray> buff;
+ std::vector<OpStatePtr> states;
+ };
+ std::mutex mutex_;
+ CachedOpConfig config_;
+ nnvm::Graph fwd_graph_;
+ nnvm::Graph grad_graph_;
+ nnvm::Graph full_graph_;
+ std::unordered_map<Context, std::vector<NDArray> > params_;
+ bool inlining_;
+ std::vector<nnvm::NodeEntry> ograd_entries_;
+ std::vector<bool> curr_grad_req_;
+ std::vector<uint32_t> bwd_in_dep_, bwd_out_dep_, bwd_ograd_dep_;
+ std::vector<uint32_t> fwd_args_idx_;
+ std::vector<uint32_t> fwd_params_idx_;
+ std::vector<uint32_t> bwd_input_eid_;
+ std::vector<bool> save_inputs_, save_outputs_;
+ };
/*! \brief whether operator recording is on. */
bool is_training() const {
return is_train_;
@@ -144,6 +222,15 @@ class Imperative {
uint32_t num_inputs, uint32_t num_outputs,
std::vector<bool> *p_save_inputs,
std::vector<bool> *p_save_outputs);
+ void RunGraph(
+ const bool retain_graph,
+ const nnvm::IndexedGraph& idx,
+ const std::vector<NDArray*> arrays,
+ size_t node_start, size_t node_end,
+ std::vector<OpReqType>&& array_reqs,
+ std::vector<uint32_t>&& ref_count,
+ std::vector<OpStatePtr> *p_states,
+ const DispatchModeVector& dispatch_modes);
/*! \brief indicate whether is training. */
#if DMLC_CXX11_THREAD_LOCAL
static thread_local bool is_train_;
@@ -160,5 +247,7 @@ class Imperative {
int backward_bulk_size_{0};
};
+using CachedOpPtr = std::shared_ptr<Imperative::CachedOp>;
+
} // namespace mxnet
#endif // MXNET_IMPERATIVE_H_
diff --git a/include/mxnet/ndarray.h b/include/mxnet/ndarray.h
index ae96fd8..e243eb7 100644
--- a/include/mxnet/ndarray.h
+++ b/include/mxnet/ndarray.h
@@ -155,14 +155,6 @@ class NDArray {
return byte_offset_ > 0 || shape() != ptr_->storage_shape;
}
- /* \brief Check whether the two arrays are the same array */
- inline bool IsSame(const NDArray& other) {
- return ptr_ == other.ptr_ &&
- shape_ == other.shape_ &&
- byte_offset_ == other.byte_offset_ &&
- dtype_ == other.dtype_;
- }
-
/*!
* \return the shape of current NDArray.
*/
diff --git a/include/mxnet/op_attr_types.h b/include/mxnet/op_attr_types.h
index f4694ef..3969d84 100644
--- a/include/mxnet/op_attr_types.h
+++ b/include/mxnet/op_attr_types.h
@@ -126,36 +126,25 @@ class OpStatePtr {
template<typename T, typename... Args>
static OpStatePtr Create(Args&&... args) {
OpStatePtr ret;
- auto state = new T(std::forward<Args>(args)...);
- auto var = Engine::Get()->NewVariable();
- ret.ptr_.reset(
- new OpState(var, state),
- [](OpState* p) {
- Engine::Get()->DeleteVariable([](RunContext s) {}, Context::CPU(), p->var);
- delete reinterpret_cast<T*>(p->state);
- delete p;
- });
+ ret.ptr_ = std::make_shared<OpState>();
+ ret.ptr_->var_ = Engine::Get()->NewVariable();
+ ret.ptr_->state_.construct<T>(std::forward<Args>(args)...);
return ret;
}
/* \brief Get engine variable associated with this state */
engine::VarHandle get_var() const {
- return ptr_->var;
+ return ptr_->var_;
}
/* \brief Get state of type T */
template<typename T>
T& get_state() const {
- return *reinterpret_cast<T*>(ptr_->state);
+ return dmlc::get<T>(ptr_->state_);
}
/* \brief clear state */
void reset() {
ptr_.reset();
}
- /* \brief checks whether the managed object is managed only by the current
- OpStatePtr instance */
- bool unique() const {
- return ptr_.unique();
- }
/* \brief Whether state is empty */
explicit operator bool() const {
return ptr_ ? true : false;
@@ -164,12 +153,16 @@ class OpStatePtr {
private:
/* \brief state structure */
struct OpState {
- engine::VarHandle var;
- void* state;
-
- OpState(engine::VarHandle var_, void* state_) : var(var_), state(state_) {}
+ OpState() {}
OpState(const OpState& other) = delete;
OpState& operator=(const OpState& other) = delete;
+
+ ~OpState() {
+ Engine::Get()->DeleteVariable([](RunContext s) {}, Context::CPU(), var_);
+ }
+
+ engine::VarHandle var_;
+ dmlc::any state_;
};
/* \brief shared pointer to state */
std::shared_ptr<OpState> ptr_;
diff --git a/python/mxnet/_ctypes/ndarray.py b/python/mxnet/_ctypes/ndarray.py
index f324545..d2cae0c 100644
--- a/python/mxnet/_ctypes/ndarray.py
+++ b/python/mxnet/_ctypes/ndarray.py
@@ -105,14 +105,28 @@ def _imperative_invoke(handle, ndargs, keys, vals, out):
class CachedOp(object):
"""Cached operator handle."""
__slots__ = ["handle"]
- def __init__(self, sym, flags=()):
+ def __init__(self, sym, flags=(), inputs=None, params=None):
self.handle = CachedOpHandle()
+ param_names = []
+ param_arrays = []
+ if inputs is None:
+ assert params is None, "When inputs is None params must also be None."
+ inputs = sym.list_inputs()
+ elif params is not None:
+ for name, arrs in params.items():
+ param_arrays.extend(arrs)
+ param_names.extend([name] * len(arrs))
check_call(_LIB.MXCreateCachedOpEx(
sym.handle,
len(flags),
c_str_array([key for key, _ in flags]),
c_str_array([str(val) for _, val in flags]),
+ len(inputs),
+ c_str_array(inputs),
+ len(param_names),
+ c_str_array(param_names),
+ c_handle_array(param_arrays),
ctypes.byref(self.handle)))
def __del__(self):
diff --git a/python/mxnet/gluon/block.py b/python/mxnet/gluon/block.py
index 293fafa..3b97c05 100644
--- a/python/mxnet/gluon/block.py
+++ b/python/mxnet/gluon/block.py
@@ -502,16 +502,8 @@ class Block(object):
----------
active : bool, default True
Whether to turn hybrid on or off.
- static_alloc : bool, default False
- Statically allocate memory to improve speed. Memory usage may increase.
- static_shape : bool, default False
- Optimize for invariant input shapes between iterations. Must also
- set static_alloc to True. Change of input shapes is still allowed
- but slower.
- forward_bulk_size : int, default 15
- Segment size of bulk execution during forward pass.
- backward_bulk_size : int, default 15
- Segment size of bulk execution during backward pass.
+ **kwargs : string
+ Additional flags for hybridized operator.
"""
for cld in self._children.values():
cld.hybridize(active, **kwargs)
@@ -704,7 +696,7 @@ class HybridBlock(Block):
self._out_format = None
self._in_format = None
self._active = False
- self._flags = []
+ self._flags = {}
def __setattr__(self, name, value):
"""Registers parameters."""
@@ -731,43 +723,39 @@ class HybridBlock(Block):
return self._cached_graph
def _build_cache(self, *args):
- data, out = self._get_graph(*args)
- data_names = {data.name : i for i, data in enumerate(data)}
- params = self.collect_params()
- input_names = out.list_inputs()
+ inputs, out = self._get_graph(*args)
+ input_names = [i.name for i in inputs]
+ params = self.collect_params()
param_names = set(params.keys())
- expected_names = set(input_names)
+ expected_names = set(out.list_inputs())
for name in expected_names:
- assert name in param_names or name in data_names, \
+ assert name in param_names or name in input_names, \
"Unknown input to HybridBlock: %s"%name
- used_data_names = [i for i in data_names if i in expected_names]
- if len(used_data_names) != len(data_names):
- unused = ', '.join(['%d-th'%i for name, i in data_names.items()
+ used_input_names = [i for i in input_names if i in expected_names]
+ if len(used_input_names) != len(input_names):
+ unused = ', '.join(['%d-th'%i for i, name in enumerate(input_names)
if name not in expected_names])
warnings.warn("The %s input to HybridBlock is not used by any "
"computation. Is this intended?"%unused, stacklevel=4)
- used_param_names = [i for i in param_names if i in expected_names]
+ used_param_names = set(i for i in param_names if i in expected_names)
if len(used_param_names) != len(param_names):
- unused = ', '.join(list(param_names - set(used_param_names)))
+ unused = ', '.join(list(param_names - used_param_names))
warnings.warn("Parameter %s is not used by any computation. "
"Is this intended?"%unused, stacklevel=4)
- data_indices = []
- param_indices = []
- self._cached_op_args = []
- for i, name in enumerate(input_names):
- if name in data_names:
- data_indices.append(i)
- self._cached_op_args.append((True, data_names[name]))
- else:
- param_indices.append(i)
- self._cached_op_args.append((False, params[name]))
- flags = [('data_indices', data_indices), ('param_indices', param_indices)] + \
- self._flags
- self._cached_op = ndarray.CachedOp(out, flags)
+ used_params = {k: params[k] for k in used_param_names}
+ try:
+ param_dict = {k: v.list_data() for k, v in used_params.items()}
+ except DeferredInitializationError:
+ self._deferred_infer_shape(*args)
+ for i in used_params.values():
+ i._finish_deferred_init()
+ param_dict = {k: v.list_data() for k, v in used_params.items()}
+
+ self._cached_op = ndarray.CachedOp(out, self._flags, input_names, param_dict)
def _deferred_infer_shape(self, *args):
try:
@@ -783,19 +771,7 @@ class HybridBlock(Block):
args, fmt = _flatten(args, "input")
assert fmt == self._in_format, "Invalid input format"
- try:
- cargs = [args[i] if is_arg else i.data()
- for is_arg, i in self._cached_op_args]
- except DeferredInitializationError:
- self._deferred_infer_shape(*args)
- cargs = []
- for is_arg, i in self._cached_op_args:
- if is_arg:
- cargs.append(args[i])
- else:
- i._finish_deferred_init()
- cargs.append(i.data())
- out = self._cached_op(*cargs)
+ out = self._cached_op(*args)
if isinstance(out, NDArray):
out = [out]
return _regroup(out, self._out_format)[0]
@@ -816,7 +792,7 @@ class HybridBlock(Block):
def hybridize(self, active=True, **kwargs):
self._active = active
- self._flags = list(kwargs.items())
+ self._flags = kwargs.items()
self._clear_cached_op()
if active and self._forward_hooks or self._forward_pre_hooks:
warnings.warn('"{}" is being hybridized while still having forward hook/pre-hook. '
diff --git a/src/c_api/c_api_ndarray.cc b/src/c_api/c_api_ndarray.cc
index 34bd4b2..9aabe04 100644
--- a/src/c_api/c_api_ndarray.cc
+++ b/src/c_api/c_api_ndarray.cc
@@ -36,7 +36,6 @@
#include "../common/utils.h"
#include "../common/exec_utils.h"
#include "../imperative/imperative_utils.h"
-#include "../imperative/cached_op.h"
using namespace mxnet;
@@ -161,8 +160,12 @@ int MXCreateCachedOp(SymbolHandle handle,
std::vector<std::string> input_names;
input_names.reserve(inputs.size());
for (const auto& i : inputs) input_names.push_back(i->attrs.name);
- *out = new CachedOpPtr(new CachedOp(
- *sym, std::vector<std::pair<std::string, std::string> >()));
+ *out = new std::shared_ptr<Imperative::CachedOp>(
+ new Imperative::CachedOp(
+ *sym,
+ std::vector<std::pair<std::string, std::string> >(),
+ input_names,
+ std::unordered_map<std::string, std::vector<NDArray> >()));
API_END();
}
@@ -170,6 +173,11 @@ int MXCreateCachedOpEx(SymbolHandle handle,
int num_flags,
const char** keys,
const char** vals,
+ int num_args,
+ const char** arg_names,
+ int num_params,
+ const char** param_names,
+ NDArrayHandle* params,
CachedOpHandle *out) {
nnvm::Symbol* sym = static_cast<nnvm::Symbol*>(handle);
@@ -178,7 +186,17 @@ int MXCreateCachedOpEx(SymbolHandle handle,
for (int i = 0; i < num_flags; ++i) {
flags.push_back({keys[i], vals[i]});
}
- *out = new CachedOpPtr(new CachedOp(*sym, flags));
+ std::vector<std::string> args;
+ for (int i = 0; i < num_args; ++i) {
+ args.push_back(arg_names[i]);
+ }
+ std::unordered_map<std::string, std::vector<NDArray> > param_dict;
+ for (int i = 0; i < num_params; ++i) {
+ param_dict[param_names[i]].emplace_back(
+ *reinterpret_cast<NDArray*>(params[i]));
+ }
+ *out = new std::shared_ptr<Imperative::CachedOp>(
+ new Imperative::CachedOp(*sym, flags, args, param_dict));
API_END();
}
diff --git a/src/engine/threaded_engine.cc b/src/engine/threaded_engine.cc
index e70cc19..dc0436e 100644
--- a/src/engine/threaded_engine.cc
+++ b/src/engine/threaded_engine.cc
@@ -278,8 +278,6 @@ void ThreadedEngine::DeleteOperator(OprHandle op) {
}
void ThreadedEngine::Push(OprHandle op, Context exec_ctx, int priority, bool profiling) {
- BulkFlush();
-
ThreadedOpr* threaded_opr = ThreadedOpr::CastFromBase(op);
OprBlock* opr_block = OprBlock::New();
opr_block->opr = threaded_opr;
@@ -325,6 +323,7 @@ void ThreadedEngine::PushAsync(AsyncFn fn, Context exec_ctx,
<< device_count_;
}
#endif
+ BulkFlush();
ThreadedOpr *opr = NewOperator(std::move(fn), const_vars, mutable_vars, prop, opr_name, wait);
opr->temporary = true;
const bool profiling = profiler_->IsProfiling(profiler::Profiler::kImperative);
diff --git a/src/executor/attach_op_execs_pass.cc b/src/executor/attach_op_execs_pass.cc
index 72919d9..697e486 100644
--- a/src/executor/attach_op_execs_pass.cc
+++ b/src/executor/attach_op_execs_pass.cc
@@ -134,10 +134,6 @@ class StatefulComputeExecutor : public StorageFallbackOpExecutor {
return state_.get_var();
}
- OpStatePtr state() const override {
- return state_;
- }
-
explicit StatefulComputeExecutor(const OpStatePtr& state,
const FStatefulCompute& fcompute,
ExecType exec_type,
@@ -146,6 +142,7 @@ class StatefulComputeExecutor : public StorageFallbackOpExecutor {
state_(state), fcompute_(fcompute), exec_type_(exec_type) {}
private:
+ friend Graph AttachOpExecs(Graph g);
OpStatePtr state_;
FStatefulCompute fcompute_;
ExecType exec_type_;
@@ -173,16 +170,13 @@ class StatefulComputeExExecutor : public OpExecutor {
return state_.get_var();
}
- OpStatePtr state() const override {
- return state_;
- }
-
explicit StatefulComputeExExecutor(const OpStatePtr& state,
const FStatefulComputeEx& fcompute,
ExecType exec_type)
: state_(state), fcompute_(fcompute), exec_type_(exec_type) {}
private:
+ friend Graph AttachOpExecs(Graph g);
OpStatePtr state_;
FStatefulComputeEx fcompute_;
ExecType exec_type_;
@@ -247,15 +241,16 @@ class FComputeExExecutor : public OpExecutor {
ExecType exec_type_;
};
-void CreateOpExecs(const Graph& g, OpExecVector* p_ret, size_t i) {
+// pass to attach operator executors
+Graph AttachOpExecs(Graph g) {
using nnvm::DTypeVector;
using nnvm::ShapeVector;
using nnvm::FMutateInputs;
- static auto& fcreate_op_state = nnvm::Op::GetAttr<FCreateOpState>("FCreateOpState");
- static auto& fmutate_inputs = nnvm::Op::GetAttr<FMutateInputs>("FMutateInputs");
- static auto& fexec_type = nnvm::Op::GetAttr<FExecType>("FExecType");
- static auto& is_layer_backward = nnvm::Op::GetAttr<bool>("TIsLayerOpBackward");
+ auto& fcreate_op_state = nnvm::Op::GetAttr<FCreateOpState>("FCreateOpState");
+ auto& fmutate_inputs = nnvm::Op::GetAttr<FMutateInputs>("FMutateInputs");
+ auto& fexec_type = nnvm::Op::GetAttr<FExecType>("FExecType");
+ auto& is_layer_backward = nnvm::Op::GetAttr<bool>("TIsLayerOpBackward");
const auto& vdtype = g.GetAttr<DTypeVector>("dtype");
const auto& vshape = g.GetAttr<ShapeVector>("shape");
@@ -264,88 +259,82 @@ void CreateOpExecs(const Graph& g, OpExecVector* p_ret, size_t i) {
// get the graph
const auto& idx = g.indexed_graph();
- OpExecVector& ret = *p_ret;
+ std::vector<std::shared_ptr<OpExecutor> > ret(idx.num_nodes());
// initialize the nodes
- const auto& inode = idx[i];
- if (inode.source->is_variable()) return;
- const nnvm::Op *op = inode.source->op();
- ExecType exec_type = ExecType::kSync;
- std::vector<uint32_t> mutate_index;
- if (fmutate_inputs.count(op)) {
- mutate_index = fmutate_inputs[op](inode.source->attrs);
- }
- if (fexec_type.count(op)) {
- exec_type = fexec_type[op](inode.source->attrs);
- }
- CHECK(dispatch_modes[i] != DispatchMode::kUndefined);
- if (fcreate_op_state.count(op)) {
- std::vector<TShape> ishape;
- std::vector<int> itype;
- for (const auto& e : inode.inputs) {
- ishape.emplace_back(vshape[idx.entry_id(e)]);
- itype.emplace_back(vdtype[idx.entry_id(e)]);
- }
-
- OpStatePtr state = fcreate_op_state[op](
- inode.source->attrs, vctx[i], ishape, itype);
- FStatefulComputeEx fcompute_ex = common::GetFCompute<FStatefulComputeEx>(
- 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<StatefulComputeExExecutor>(state, fcompute_ex, exec_type);
- } else {
- FStatefulCompute fcompute = common::GetFCompute<FStatefulCompute>(
- 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<StatefulComputeExecutor>(state, fcompute,
- exec_type, mutate_index);
+ for (size_t i = 0; i < idx.num_nodes(); ++i) {
+ const auto& inode = idx[i];
+ if (inode.source->is_variable()) continue;
+ const nnvm::Op *op = inode.source->op();
+ ExecType exec_type = ExecType::kSync;
+ std::vector<uint32_t> mutate_index;
+ if (fmutate_inputs.count(op)) {
+ mutate_index = fmutate_inputs[op](inode.source->attrs);
}
- } else if (is_layer_backward.get(op, false)) {
- CHECK_GE(inode.control_deps.size(), 1);
- uint32_t fwd_id = inode.control_deps[0];
- CHECK(vctx[fwd_id] == vctx[i]);
- CHECK(ret[fwd_id] != nullptr);
- FStatefulComputeEx fcompute_ex = common::GetFCompute<FStatefulComputeEx>(
- 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<StatefulComputeExExecutor>(
- ret[fwd_id].get()->state(), fcompute_ex, exec_type);
- } else {
- FStatefulCompute fcompute = common::GetFCompute<FStatefulCompute>(
- 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<StatefulComputeExecutor>(
- ret[fwd_id].get()->state(), fcompute, exec_type, mutate_index);
+ if (fexec_type.count(op)) {
+ exec_type = fexec_type[op](inode.source->attrs);
}
- } else {
- FCompute fcompute = common::GetFCompute<FCompute>(op, "FCompute", vctx[i]);
- FComputeEx fcomp_ex = common::GetFCompute<FComputeEx>(op, "FComputeEx", vctx[i]);
- if (fcomp_ex != nullptr && dispatch_modes[i] == DispatchMode::kFComputeEx) {
- ret[i] = std::make_shared<FComputeExExecutor>(
- inode.source->attrs, fcomp_ex, exec_type);
- } else if (fcompute != nullptr) {
- ret[i] = std::make_shared<FComputeExecutor>(
- inode.source->attrs, fcompute, exec_type, mutate_index);
+ CHECK(dispatch_modes[i] != DispatchMode::kUndefined);
+ if (fcreate_op_state.count(op)) {
+ std::vector<TShape> ishape;
+ std::vector<int> itype;
+ for (const auto& e : inode.inputs) {
+ ishape.emplace_back(vshape[idx.entry_id(e)]);
+ itype.emplace_back(vdtype[idx.entry_id(e)]);
+ }
+
+ OpStatePtr state = fcreate_op_state[op](
+ inode.source->attrs, vctx[i], ishape, itype);
+ FStatefulComputeEx fcompute_ex = common::GetFCompute<FStatefulComputeEx>(
+ 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<StatefulComputeExExecutor>(state, fcompute_ex, exec_type);
+ } else {
+ FStatefulCompute fcompute = common::GetFCompute<FStatefulCompute>(
+ 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<StatefulComputeExecutor>(state, fcompute,
+ exec_type, mutate_index);
+ }
+ } else if (is_layer_backward.get(op, false)) {
+ CHECK_GE(inode.control_deps.size(), 1);
+ uint32_t fwd_id = inode.control_deps[0];
+ CHECK(vctx[fwd_id] == vctx[i]);
+ CHECK(ret[fwd_id] != nullptr);
+ FStatefulComputeEx fcompute_ex = common::GetFCompute<FStatefulComputeEx>(
+ 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<StatefulComputeExExecutor>(
+ dynamic_cast<StatefulComputeExExecutor*>(ret[fwd_id].get())->state_,
+ fcompute_ex, exec_type);
+ } else {
+ FStatefulCompute fcompute = common::GetFCompute<FStatefulCompute>(
+ 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<StatefulComputeExecutor>(
+ dynamic_cast<StatefulComputeExecutor*>(ret[fwd_id].get())->state_,
+ fcompute, exec_type, mutate_index);
+ }
} else {
- LOG(INFO) << "Neither FCompute nor FComputeEx registered " << op->name;
+ FCompute fcompute = common::GetFCompute<FCompute>(op, "FCompute", vctx[i]);
+ FComputeEx fcomp_ex = common::GetFCompute<FComputeEx>(op, "FComputeEx", vctx[i]);
+ if (fcomp_ex != nullptr && dispatch_modes[i] == DispatchMode::kFComputeEx) {
+ ret[i] = std::make_shared<FComputeExExecutor>(
+ inode.source->attrs, fcomp_ex, exec_type);
+ } else if (fcompute != nullptr) {
+ ret[i] = std::make_shared<FComputeExecutor>(
+ inode.source->attrs, fcompute, exec_type, mutate_index);
+ } else {
+ LOG(INFO) << "Neither FCompute nor FComputeEx registered " << op->name;
+ }
}
}
-}
-
-
-// pass to attach operator executors
-Graph AttachOpExecs(Graph g) {
- const auto& idx = g.indexed_graph();
- OpExecVector ret(idx.num_nodes());
- for (size_t i = 0; i < idx.num_nodes(); ++i) {
- CreateOpExecs(g, &ret, i);
- }
g.attrs["op_execs"] = std::make_shared<nnvm::any>(ret);
return g;
}
diff --git a/src/executor/attach_op_resource_pass.cc b/src/executor/attach_op_resource_pass.cc
index 56122cd..6818662 100644
--- a/src/executor/attach_op_resource_pass.cc
+++ b/src/executor/attach_op_resource_pass.cc
@@ -30,15 +30,12 @@
namespace mxnet {
namespace exec {
-void AttachOpResources(
- const Graph& g,
- const OpExecVector& op_execs,
- size_t start_nid,
- size_t end_nid) {
+Graph AttachOpResources(Graph g) {
static auto& fresource =
nnvm::Op::GetAttr<FResourceRequest>("FResourceRequest");
static auto& fresource_ex =
nnvm::Op::GetAttr<FResourceRequestEx>("FResourceRequestEx");
+ auto& op_execs = nnvm::get<OpExecVector>(*g.attrs.at("op_execs"));
const auto& vctx = g.GetAttr<ContextVector>("context");
const auto& vdispatch = g.GetAttr<DispatchModeVector>("dispatch_mode");
const auto& dev_masks = g.GetAttr<DevMaskVector>("dev_mask");
@@ -46,7 +43,7 @@ void AttachOpResources(
// Use global resource pool for each executor for now.
std::map<Context, Resource> cached_temp;
// Resource allocation
- for (uint32_t nid = start_nid; nid < end_nid; ++nid) {
+ for (uint32_t nid = 0; nid < idx.num_nodes(); ++nid) {
const auto& inode = idx[nid];
if (inode.source->is_variable()) continue;
const Context &ctx = vctx[nid];
@@ -87,12 +84,7 @@ void AttachOpResources(
requested.push_back(ResourceManager::Get()->Request(ctx, ResourceRequest::kTempSpace));
}
}
+ return g;
}
-
-void AttachOpResources(const Graph& g) {
- const auto& op_execs = g.GetAttr<OpExecVector>("op_execs");
- AttachOpResources(g, op_execs, 0, g.indexed_graph().num_nodes());
-}
-
} // namespace exec
} // namespace mxnet
diff --git a/src/executor/exec_pass.h b/src/executor/exec_pass.h
index 26a2491..99b1b16 100644
--- a/src/executor/exec_pass.h
+++ b/src/executor/exec_pass.h
@@ -82,10 +82,6 @@ class OpExecutor {
virtual engine::VarHandle var() const {
return nullptr;
}
- /*! \return return operator state */
- virtual OpStatePtr state() const {
- return OpStatePtr();
- }
};
/*!
@@ -107,14 +103,6 @@ using ContextVector = std::vector<Context>;
using DevMaskVector = std::vector<int>;
/*!
- * \brief create OpExecutor for a node in graph
- *
- * \param g input graph
- * \param p_ret OpExecVector for input and output
- * \param i the id of the node
- */
-void CreateOpExecs(const Graph& g, OpExecVector* p_ret, size_t i);
-/*!
* \brief Attach OpExecutor to the graph attributes.
*
* \param g input graph
@@ -127,20 +115,12 @@ Graph AttachOpExecs(Graph g);
* \brief Attach Resource to the OpExecVector of the graph.
*
* \param g input graph need to contain op_exec attribute.
- */
-void AttachOpResources(const Graph& g);
-/*!
- * \brief Attach Resource to the OpExecVector
*
- * \param g input graph
- * \param op_execs OpExecutor vector
- * \param start_nid starting node id
- * \param end_nid end node id
+ * \return graph with new attribute "op_exec" of type OpExecVector
+ * The fields on the OpExecVector are not yet been setup.
*/
-void AttachOpResources(const Graph& g,
- const OpExecVector& op_execs,
- size_t start_nid,
- size_t end_nid);
+Graph AttachOpResources(Graph g);
+
/*!
* \brief Discover chance of inplace addto operators.
* i.e. z = plus(z, source_op), and encourage it to become z += source_op.
diff --git a/src/executor/graph_executor.cc b/src/executor/graph_executor.cc
index 831b5f9..e28867d 100644
--- a/src/executor/graph_executor.cc
+++ b/src/executor/graph_executor.cc
@@ -912,7 +912,7 @@ void GraphExecutor::FinishInitGraph(nnvm::Symbol symbol,
}
g = AttachOpExecs(g);
- AttachOpResources(g);
+ g = AttachOpResources(g);
graph_ = std::move(g);
if (shared_exec != nullptr) {
diff --git a/src/imperative/cached_op.cc b/src/imperative/cached_op.cc
index b17fae4..140b5a5 100644
--- a/src/imperative/cached_op.cc
+++ b/src/imperative/cached_op.cc
@@ -19,78 +19,16 @@
#include <unordered_set>
#include <iostream>
#include "./imperative_utils.h"
-#include "./cached_op.h"
-#include "../executor/exec_pass.h"
-#include "../profiler/profiler.h"
-
namespace mxnet {
DMLC_REGISTER_PARAMETER(CachedOpConfig);
-struct CachedOp::GraphInfo {
- nnvm::Graph fwd_graph;
- nnvm::Graph full_graph;
- std::vector<OpReqType> bwd_output_reqs;
- std::vector<uint32_t> bwd_input_eid;
-};
-
-struct CachedOp::DynamicRuntime {
- GraphInfo info;
- std::vector<NDArray> buff;
- std::vector<OpStatePtr> op_states;
-};
-
-struct CachedOp::CachedOpState {
- CachedOpState(const Context& context_,
- const nnvm::Graph& fwd_graph_,
- const nnvm::Graph& full_graph_) {
- context = context_;
- info.fwd_graph = fwd_graph_;
- info.full_graph = full_graph_;
-
- size_t max_nodes = info.full_graph.indexed_graph().num_nodes();
- size_t max_entries = info.full_graph.indexed_graph().num_node_entries();
- info.fwd_graph.attrs["context"] = std::make_shared<dmlc::any>(
- std::vector<Context>(info.fwd_graph.indexed_graph().num_nodes(), context));
- info.full_graph.attrs["context"] = std::make_shared<dmlc::any>(
- std::vector<Context>(max_nodes, context));
-
- buff.resize(max_entries);
- arrays.resize(max_entries);
- array_reqs.resize(max_entries);
- dynamic_entries.resize(max_entries, false);
- op_states.resize(max_nodes);
- execs.resize(max_nodes);
- opr_segs.resize(max_nodes);
- }
-
- std::mutex mutex;
- Context context;
- GraphInfo info;
-
- bool recording = false;
- bool fwd_alloc = false;
- bool bwd_alloc = false;
- bool fwd_exec_init = false;
- bool bwd_exec_init = false;
-
- std::vector<NDArray> buff;
- std::vector<NDArray*> arrays;
- std::vector<OpReqType> array_reqs;
-
- std::vector<OpStatePtr> op_states;
- std::vector<std::shared_ptr<exec::OpExecutor> > execs;
- std::vector<imperative::EngineOprSeg> opr_segs;
-
- std::vector<bool> dynamic_entries;
- std::multimap<size_t, NDArray> fwd_reuse_pool;
- std::multimap<size_t, NDArray> bwd_reuse_pool;
-};
-
-CachedOp::CachedOp(
+Imperative::CachedOp::CachedOp(
const nnvm::Symbol& sym,
- const std::vector<std::pair<std::string, std::string> >& flags) {
+ const std::vector<std::pair<std::string, std::string> >& flags,
+ const std::vector<std::string> arg_names,
+ const std::unordered_map<std::string, std::vector<NDArray> >& params) {
using namespace nnvm;
using namespace imperative;
static const std::vector<const Op*> zero_ops{Op::Get("zeros_like"), Op::Get("_zeros")};
@@ -130,22 +68,34 @@ CachedOp::CachedOp(
fwd_graph_.attrs["forward_ref_count"] =
std::make_shared<dmlc::any>(std::move(ref_count));
- inlining_ = !config_.static_alloc &&
- (idx.num_nodes() - idx.input_nodes().size()) <= config_.inline_limit;
+ inlining_ = (idx.num_nodes() - idx.input_nodes().size()) <= config_.inline_limit;
}
// Set params
{
const auto& idx = fwd_graph_.indexed_graph();
- if (config_.data_indices.ndim() || config_.param_indices.ndim()) {
- CHECK_EQ(config_.data_indices.ndim() + config_.param_indices.ndim(),
- idx.input_nodes().size());
- } else {
- std::vector<uint32_t> tmp;
- for (size_t i = 0; i < idx.input_nodes().size(); ++i) {
- tmp.push_back(i);
+ std::unordered_map<std::string, size_t> arg_name_to_id;
+ for (size_t i = 0; i < idx.input_nodes().size(); ++i) {
+ const auto& name = idx[idx.input_nodes()[i]].source->attrs.name;
+ auto iter = params.find(name);
+ if (iter == params.end()) {
+ arg_name_to_id[name] = i;
+ continue;
+ }
+ fwd_params_idx_.push_back(i);
+ for (const auto& param : iter->second) {
+ params_[param.ctx()].emplace_back(param);
}
- config_.data_indices.assign(tmp.begin(), tmp.end());
+ }
+
+ CHECK_EQ(arg_name_to_id.size(), arg_names.size())
+ << "CachedOp expects " << arg_name_to_id.size()
+ << " inputs, given " << arg_names.size();
+
+ for (const auto& name : arg_names) {
+ auto iter = arg_name_to_id.find(name);
+ CHECK(iter != arg_name_to_id.end()) << "Unexpected input name " << name;
+ fwd_args_idx_.push_back(iter->second);
}
}
@@ -157,14 +107,9 @@ CachedOp::CachedOp(
}
std::vector<NodeEntry> xs;
- const auto& idx = fwd_graph_.indexed_graph();
- for (size_t i = 0; i < idx.input_nodes().size(); ++i) {
- auto nid = idx.input_nodes()[i];
- if (idx.mutable_input_nodes().count(nid)) continue;
- fwd_input_to_grad_output_[i] = xs.size();
- xs.emplace_back(NodeEntry{idx[nid].weak_ref.lock(), 0, 0});
- }
-
+ std::vector<NodePtr> args = sym.ListInputs(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.";
@@ -180,7 +125,7 @@ CachedOp::CachedOp(
size_t num_forward_entries = fwd_graph_.indexed_graph().num_node_entries();
full_graph_.outputs = fwd_graph_.outputs;
- bwd_output_reqs_ = std::vector<OpReqType>(grad_graph_.outputs.size(), kWriteTo);
+ curr_grad_req_ = std::vector<bool>(grad_graph_.outputs.size(), true);
for (const auto& i : grad_graph_.outputs) full_graph_.outputs.emplace_back(i);
const auto& idx = full_graph_.indexed_graph();
@@ -224,10 +169,7 @@ CachedOp::CachedOp(
}
}
-CachedOp::~CachedOp() {
-}
-
-std::vector<nnvm::NodeEntry> CachedOp::Gradient(
+std::vector<nnvm::NodeEntry> Imperative::CachedOp::Gradient(
const nnvm::NodePtr& node,
const std::vector<nnvm::NodeEntry>& ograds) {
using namespace nnvm;
@@ -264,15 +206,13 @@ std::vector<nnvm::NodeEntry> CachedOp::Gradient(
return ret;
}
-
-bool CachedOp::SetForwardGraph(
- GraphInfo* info,
- const bool recording,
- const std::vector<NDArray*>& inputs) {
+nnvm::Graph Imperative::CachedOp::GetForwardGraph(
+ const bool recording, const std::vector<NDArray*>& inputs) {
using namespace nnvm;
using namespace imperative;
+ std::lock_guard<std::mutex> lock(mutex_);
CHECK_EQ(inputs.size(), num_inputs());
- nnvm::Graph& g = info->fwd_graph;
+ nnvm::Graph& g = fwd_graph_;
ShapeVector shape_inputs;
DTypeVector dtype_inputs;
@@ -297,22 +237,18 @@ bool CachedOp::SetForwardGraph(
g.attrs.erase("forward_mem_plan");
g.attrs.erase("full_mem_plan");
} else if (g.attrs.count(recording ? "full_mem_plan" : "forward_mem_plan")) {
- return true;
+ return g;
}
const auto& idx = g.indexed_graph();
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<StorageTypeVector>("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;
- }
- for (const auto i : idx.input_nodes()) {
- storage[idx.entry_id(i, 0)] = exec::kExternalStorageID;
- }
- for (size_t i = 0; i < idx.outputs().size(); ++i) {
- storage[idx.entry_id(idx.outputs()[i])] = exec::kExternalStorageID;
+ if (stypes[i] != kDefaultStorage)
+ storage[i] = exec::kDynamicStorageID;
}
auto mem_plan = PlanMemory(
@@ -321,50 +257,51 @@ bool CachedOp::SetForwardGraph(
g.attrs[recording ? "full_mem_plan" : "forward_mem_plan"] =
std::make_shared<dmlc::any>(std::move(mem_plan));
- return false;
+ return g;
}
-bool CachedOp::SetBackwardGraph(
- GraphInfo* info,
+nnvm::Graph Imperative::CachedOp::GetBackwardGraph(
+ const OpStatePtr& op_state,
const std::vector<OpReqType>& reqs,
- const std::vector<NDArray*>& inputs,
- bool detect_inplace_addto) {
+ const std::vector<NDArray*>& inputs) {
using namespace nnvm;
using namespace imperative;
std::lock_guard<std::mutex> lock(mutex_);
- Context default_ctx = inputs[0]->ctx();
- nnvm::Graph& g = info->full_graph;
-
- if (info->bwd_output_reqs != reqs) {
- info->bwd_output_reqs = reqs;
- info->bwd_input_eid.clear();
+ nnvm::Graph& g = full_graph_;
+ auto& state = op_state.get_state<CachedOpState>();
+ bool req_match = true;
+ for (size_t i = 0; i < reqs.size(); ++i) {
+ if (curr_grad_req_[i] != (reqs[i] != kNullOp)) {
+ curr_grad_req_[i] = reqs[i] != kNullOp;
+ req_match = false;
+ }
+ }
+ if (!req_match) {
g = nnvm::Graph();
g.outputs = fwd_graph_.outputs;
for (size_t i = 0; i < grad_graph_.outputs.size(); ++i) {
- if (info->bwd_output_reqs[i] == kNullOp) continue;
- g.outputs.emplace_back(grad_graph_.outputs[i]);
+ if (curr_grad_req_[i]) g.outputs.emplace_back(grad_graph_.outputs[i]);
}
- g.attrs["context"] = std::make_shared<dmlc::any>(
- std::vector<Context>(g.indexed_graph().num_nodes(), default_ctx));
+ bwd_input_eid_.clear();
}
const auto& idx = g.indexed_graph();
- if (info->bwd_input_eid.size() != inputs.size()) {
- info->bwd_input_eid.clear();
+ if (bwd_input_eid_.size() != inputs.size()) {
+ bwd_input_eid_.clear();
for (const auto& i : bwd_ograd_dep_) {
auto eid = idx.entry_id(ograd_entries_[i]);
- info->bwd_input_eid.push_back(eid);
+ bwd_input_eid_.push_back(eid);
}
for (const auto& i : bwd_in_dep_) {
auto eid = idx.entry_id(idx.input_nodes()[i], 0);
- info->bwd_input_eid.push_back(eid);
+ bwd_input_eid_.push_back(eid);
}
for (const auto& i : bwd_out_dep_) {
auto eid = idx.entry_id(idx.outputs()[i]);
- info->bwd_input_eid.push_back(eid);
+ bwd_input_eid_.push_back(eid);
}
- CHECK_EQ(inputs.size(), info->bwd_input_eid.size());
+ CHECK_EQ(inputs.size(), bwd_input_eid_.size());
}
size_t num_forward_nodes = fwd_graph_.indexed_graph().num_nodes();
@@ -375,22 +312,25 @@ bool CachedOp::SetBackwardGraph(
for (size_t i = num_forward_nodes; i < idx.num_nodes(); ++i) {
for (const auto& j : idx[i].inputs) ++ref_count[idx.entry_id(j)];
}
- for (size_t i = 0; i < inputs.size(); ++i) ++ref_count[info->bwd_input_eid[i]];
+ for (size_t i = 0; i < inputs.size(); ++i) ++ref_count[bwd_input_eid_[i]];
for (const auto& i : idx.outputs()) ++ref_count[idx.entry_id(i)];
g.attrs["backward_ref_count"] = std::make_shared<dmlc::any>(std::move(ref_count));
}
- auto shapes = info->fwd_graph.GetAttr<ShapeVector>("shape");
- shapes.resize(idx.num_node_entries(), TShape());
- auto dtypes = info->fwd_graph.GetAttr<DTypeVector>("dtype");
- dtypes.resize(idx.num_node_entries(), -1);
- auto stypes = info->fwd_graph.GetAttr<StorageTypeVector>("storage_type");
- stypes.resize(idx.num_node_entries(), -1);
+ ShapeVector shapes(idx.num_node_entries(), TShape());
+ DTypeVector dtypes(idx.num_node_entries(), -1);
+ StorageTypeVector stypes(idx.num_node_entries(), -1);
+
+ for (size_t i = 0; i < num_forward_entries; ++i) {
+ shapes[i] = state.buff[i].shape();
+ dtypes[i] = state.buff[i].dtype();
+ stypes[i] = state.buff[i].storage_type();
+ }
for (size_t i = 0; i < inputs.size(); ++i) {
- shapes[info->bwd_input_eid[i]] = inputs[i]->shape();
- dtypes[info->bwd_input_eid[i]] = inputs[i]->dtype();
- stypes[info->bwd_input_eid[i]] = inputs[i]->storage_type();
+ shapes[bwd_input_eid_[i]] = inputs[i]->shape();
+ dtypes[bwd_input_eid_[i]] = inputs[i]->dtype();
+ stypes[bwd_input_eid_[i]] = inputs[i]->storage_type();
}
std::pair<uint32_t, uint32_t> node_range, entry_range;
@@ -402,353 +342,79 @@ bool CachedOp::SetBackwardGraph(
node_range, entry_range);
match &= CheckAndInferType(&g, std::move(dtypes), false,
node_range, entry_range);
- exec::DevMaskVector dev_mask(idx.num_nodes(), default_ctx.dev_mask());
+ exec::DevMaskVector dev_mask(idx.num_nodes(), inputs[0]->ctx().dev_mask());
match &= CheckAndInferStorageType(&g, std::move(dev_mask), std::move(stypes),
false, node_range, entry_range);
if (!match) {
g.attrs.erase("backward_mem_plan");
} else if (g.attrs.count("backward_mem_plan")) {
- return true;
+ return g;
}
StorageVector storage(idx.num_node_entries(), exec::kBadStorageID);
- const auto& bwd_stypes = g.GetAttr<StorageTypeVector>("storage_type");
- for (size_t i = 0; i < bwd_stypes.size(); i++) {
- if (bwd_stypes[i] != kDefaultStorage) storage[i] = exec::kDynamicStorageID;
- }
for (size_t i = 0; i < num_forward_entries; ++i) storage[i] = exec::kExternalStorageID;
for (const auto i : idx.input_nodes()) storage[idx.entry_id(i, 0)] = exec::kExternalStorageID;
for (const auto i : idx.outputs()) storage[idx.entry_id(i)] = exec::kExternalStorageID;
+ for (size_t i = 0; i < stypes.size(); i++) {
+ if (stypes[i] != kDefaultStorage)
+ storage[i] = exec::kDynamicStorageID;
+ }
auto mem_plan = PlanMemory(
&g, std::move(storage), g.GetAttr<std::vector<uint32_t> >("backward_ref_count"),
- {num_forward_nodes, idx.num_nodes()},
- {num_forward_entries, idx.num_node_entries()},
- detect_inplace_addto);
+ {num_forward_nodes, idx.num_nodes()}, {num_forward_entries, idx.num_node_entries()});
g.attrs["backward_mem_plan"] = std::make_shared<dmlc::any>(std::move(mem_plan));
- return false;
-}
-
-OpStatePtr CachedOp::GetCachedOpState(
- const Context& ctx) {
- std::lock_guard<std::mutex> lock(mutex_);
- for (const auto& i : cached_op_states_[ctx]) {
- // only create one state per device when not using static memory
- if (!config_.static_alloc || i.unique()) {
- return i;
- }
- }
- auto state_ptr = OpStatePtr::Create<CachedOpState>(ctx, fwd_graph_, full_graph_);
-
- cached_op_states_[ctx].push_back(state_ptr);
- return state_ptr;
-}
-
-void CachedOp::StaticAllocMemory(
- const OpStatePtr& state_ptr,
- bool recording,
- bool keep_fwd) {
- using namespace nnvm;
- using namespace imperative;
-
- auto& state = state_ptr.get_state<CachedOpState>();
- const auto& default_ctx = state.context;
- nnvm::Graph& g = keep_fwd ? state.info.full_graph : state.info.fwd_graph;
- const auto& idx = g.indexed_graph();
- const auto& vstorage_inplace = g.GetAttr<std::vector<int> >("storage_inplace_index");
- const auto& mem_plan = g.GetAttr<MemoryPlanVector>(
- keep_fwd ? "backward_mem_plan" : (recording ? "full_mem_plan" : "forward_mem_plan"));
- std::vector<int> addto_entry;
- if (g.attrs.count("addto_entry")) {
- addto_entry = g.GetAttr<std::vector<int> >("addto_entry");
- }
- size_t start_eid =
- keep_fwd ? state.info.fwd_graph.indexed_graph().num_node_entries() : 0;
- size_t end_eid = idx.num_node_entries();
-
- if (!keep_fwd) state.fwd_alloc = false;
- state.bwd_alloc = false;
- for (size_t i = start_eid; i < state.buff.size(); ++i) {
- state.buff[i] = NDArray();
- state.arrays[i] = &state.buff[i];
- state.array_reqs[i] = kNullOp;
- state.dynamic_entries[i] = false;
- }
-
- for (auto i : idx.input_nodes()) {
- auto eid = idx.entry_id(i, 0);
- if (eid >= start_eid) state.dynamic_entries[eid] = true;
- }
- for (auto i : idx.outputs()) {
- auto eid = idx.entry_id(i);
- if (eid >= start_eid) state.dynamic_entries[eid] = true;
- }
-
- for (size_t i = start_eid; i < end_eid; ++i) {
- if (addto_entry.size() && addto_entry[i]) {
- state.array_reqs[i] = kAddTo;
- } else if (vstorage_inplace[i] >= 0) {
- state.array_reqs[i] = kWriteInplace;
- } else if (vstorage_inplace[i] == -2) {
- // -2 indicate that the entry is never referenced.
- state.array_reqs[i] = kNullOp;
- } else {
- state.array_reqs[i] = kWriteTo;
- }
- }
-
- auto& reuse_pool = keep_fwd ? state.bwd_reuse_pool : state.fwd_reuse_pool;
- reuse_pool = imperative::AllocateMemory(
- g, idx, default_ctx, start_eid, end_eid, mem_plan,
- state.arrays, &state.array_reqs, std::move(reuse_pool));
-
- state.recording = recording;
- if (keep_fwd) {
- state.bwd_alloc = true;
- } else {
- state.fwd_alloc = true;
- }
+ return g;
}
-void CachedOp::StaticInitExec(
- const OpStatePtr& state_ptr,
- bool recording,
- bool keep_fwd) {
+void Imperative::CachedOp::Forward(
+ const std::shared_ptr<CachedOp>& op_ptr,
+ const std::vector<NDArray*>& args,
+ const std::vector<NDArray*>& outputs) {
using namespace nnvm;
using namespace imperative;
+ static const auto cached_op = nnvm::Op::Get("_CachedOp");
- auto& state = state_ptr.get_state<CachedOpState>();
- const auto& default_ctx = state.context;
- nnvm::Graph& g = keep_fwd ? state.info.full_graph : state.info.fwd_graph;
- const auto& idx = g.indexed_graph();
- std::vector<int> skip_plus_node;
- if (g.attrs.count("skip_plus_node")) {
- skip_plus_node = g.GetAttr<std::vector<int> >("skip_plus_node");
- }
- size_t start_nid =
- keep_fwd ? state.info.fwd_graph.indexed_graph().num_nodes() : 0;
- size_t end_nid = idx.num_nodes();
-
- if (!keep_fwd) state.fwd_exec_init = false;
- state.bwd_exec_init = false;
-
- for (size_t i = start_nid; i < state.execs.size(); ++i) {
- state.execs[i].reset();
- state.opr_segs[i] = EngineOprSeg();
- }
-
- if (!config_.static_shape) {
- for (size_t i = start_nid; i < end_nid; ++i) {
- state.opr_segs[i].next_nid = i + 1;
- state.opr_segs[i].skip = skip_plus_node.size() && skip_plus_node[i];
- }
- } else {
- for (size_t i = start_nid; i < end_nid; ++i) {
- exec::CreateOpExecs(g, &state.execs, i);
- }
- exec::AttachOpResources(g, state.execs, start_nid, end_nid);
-
- for (size_t i = start_nid; i < end_nid; ++i) {
- bool skip = idx[i].source->is_variable();
- for (size_t j = 0; !skip && j < idx[i].inputs.size(); ++j) {
- skip = state.dynamic_entries[idx.entry_id(idx[i].inputs[j])];
- }
- for (size_t j = 0; !skip && j < idx[i].source->num_outputs(); ++j) {
- skip = state.dynamic_entries[idx.entry_id(i, j)];
- }
- if (skip) continue;
- SetupOpExec(g, i, state.execs[i], state.arrays, state.array_reqs);
- }
+ CHECK_EQ(args.size(), fwd_args_idx_.size())
+ << "CachedOp requires " << fwd_args_idx_.size()
+ << " inputs but got " << args.size();
- size_t bulk_size = idx.num_nodes();
- std::unordered_set<uint32_t> excludes;
- if (recording || keep_fwd) {
- bulk_size = keep_fwd ? config_.backward_bulk_size : config_.forward_bulk_size;
- for (const auto& i : idx.outputs()) excludes.insert(idx.entry_id(i));
- for (const auto& i : idx.input_nodes()) excludes.insert(idx.entry_id(i, 0));
- }
+ Context default_ctx = args[0]->ctx();
- CreateEngineOpSeg(idx, default_ctx, start_nid, end_nid, bulk_size, excludes,
- state.execs, skip_plus_node, &state.opr_segs);
- }
- if (keep_fwd) {
- state.bwd_exec_init = true;
- } else {
- state.fwd_exec_init = true;
+ std::vector<NDArray*> inputs(num_inputs());
+ for (index_t i = 0; i < fwd_args_idx_.size(); ++i) {
+ inputs[fwd_args_idx_[i]] = args[i];
}
-}
-
-void CachedOp::StaticRunOps(
- const Context& default_ctx,
- const nnvm::Graph& g,
- const OpStatePtr& state_ptr,
- size_t start_nid,
- size_t end_nid) {
- static auto& createop = nnvm::Op::GetAttr<FCreateOpState>("FCreateOpState");
- static auto& is_layer_backward = Op::GetAttr<bool>("TIsLayerOpBackward");
-
- bool profiling = profiler::Profiler::Get()->GetState() == profiler::Profiler::kRunning;
- bool is_training = Imperative::Get()->is_training();
- auto& state = state_ptr.get_state<CachedOpState>();
- const auto& idx = g.indexed_graph();
- const auto& dispatch_modes = g.GetAttr<DispatchModeVector>("dispatch_mode");
- const auto& op_execs = state.execs;
-
- std::vector<NDArray*> ndinputs, ndoutputs;
- nnvm::ShapeVector arg_shapes;
- nnvm::DTypeVector arg_dtypes;
- std::vector<OpReqType> req;
+ if (fwd_params_idx_.size()) {
+ CHECK(params_.find(default_ctx) != params_.end())
+ << "CachedOp is not initialized on context " << default_ctx;
- for (size_t i = start_nid; config_.static_shape && i < end_nid; ++i) {
- if (op_execs[i]) op_execs[i]->op_ctx.is_train = is_training;
- }
-
- for (size_t i = start_nid; i < end_nid; i = state.opr_segs[i].next_nid) {
- const auto& opr_seg = state.opr_segs[i];
- if (opr_seg.skip) continue;
- if (opr_seg.opr != nullptr) {
- Engine::Get()->Push(opr_seg.opr.get(), default_ctx, 0, profiling);
- } else {
- const nnvm::IndexedGraph::Node& node = idx[i];
- if (node.source->is_variable()) continue;
- auto num_outputs = node.source->num_outputs();
- ndinputs.clear();
- ndinputs.reserve(node.inputs.size());
- for (const auto& j : node.inputs) {
- ndinputs.emplace_back(state.arrays[idx.entry_id(j)]);
- CHECK(!ndinputs.back()->is_none());
- }
- 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(state.arrays[eid]);
- req.push_back(state.array_reqs[eid]);
- CHECK(req.back() == kNullOp || !ndoutputs.back()->is_none());
- }
- 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());
- }
- state.op_states[i] = createop[node.source->op()](
- node.source->attrs, default_ctx, arg_shapes, arg_dtypes);
- Imperative::Get()->InvokeOp(
- default_ctx, node.source->attrs, ndinputs, ndoutputs, req,
- dispatch_mode, state.op_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::Get()->InvokeOp(
- default_ctx, node.source->attrs, ndinputs, ndoutputs,
- req, dispatch_mode, state.op_states[fwd_node_id]);
- } else {
- Imperative::Get()->InvokeOp(
- default_ctx, node.source->attrs, ndinputs, ndoutputs, req,
- dispatch_mode);
- }
+ for (size_t i = 0; i < fwd_params_idx_.size(); ++i) {
+ inputs[fwd_params_idx_[i]] = ¶ms_[default_ctx][i];
}
}
-}
-
-OpStatePtr CachedOp::StaticForward(
- const Context& default_ctx,
- const std::vector<NDArray*>& inputs,
- const std::vector<NDArray*>& outputs) {
- using namespace nnvm;
- using namespace imperative;
+ // Initialize
bool recording = Imperative::Get()->is_recording();
- auto state_ptr = GetCachedOpState(default_ctx);
- auto& state = state_ptr.get_state<CachedOpState>();
- std::lock_guard<std::mutex> lock(state.mutex);
-
- bool match = SetForwardGraph(&state.info, recording, inputs);
- match = match && state.recording == recording;
-
- nnvm::Graph& g = state.info.fwd_graph;
+ nnvm::Graph g = GetForwardGraph(recording, inputs);
const auto& idx = g.indexed_graph();
- if (!state.fwd_alloc || !match) {
- StaticAllocMemory(state_ptr, recording, false);
- }
-
- if (config_.static_shape) {
- for (auto i : config_.param_indices) {
- auto nid = idx.input_nodes()[i];
- if (!state.arrays[idx.entry_id(nid, 0)]->IsSame(*inputs[i])) {
- match = false;
- auto ptr = &state.buff[idx.entry_id(nid, 0)];
- CHECK_EQ(state.arrays[idx.entry_id(nid, 0)], ptr);
- *state.arrays[idx.entry_id(nid, 0)] = *inputs[i];
- state.dynamic_entries[idx.entry_id(nid, 0)] = false;
- }
- }
- for (auto i : config_.data_indices) {
- auto eid = idx.entry_id(idx.input_nodes()[i], 0);
- state.arrays[eid] = inputs[i];
- }
- } else {
- for (size_t i = 0; i < num_inputs(); ++i) {
- auto nid = idx.input_nodes()[i];
- state.arrays[idx.entry_id(nid, 0)] = inputs[i];
- }
- }
-
- if (!state.fwd_exec_init || !match) {
- StaticInitExec(state_ptr, recording, false);
- }
-
- const auto& dtypes = g.GetAttr<DTypeVector>("dtype");
- const auto& shapes = g.GetAttr<ShapeVector>("shape");
- const auto& stypes = g.GetAttr<StorageTypeVector>("storage_type");
+ size_t num_inputs = idx.input_nodes().size();
- for (size_t i = 0; i < outputs.size(); ++i) {
- auto eid = idx.entry_id(idx.outputs()[i]);
- state.arrays[eid] = outputs[i];
- if (!outputs[i]->is_none()) continue;
- *outputs[i] = NDArray(static_cast<NDArrayStorageType>(stypes[eid]),
- shapes[eid], default_ctx, true, dtypes[eid]);
+ for (size_t i = 0; i < inputs.size(); ++i) {
+ CHECK_EQ(inputs[i]->ctx(), default_ctx)
+ << "CachedOp 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 "
+ << inputs[i]->ctx();
}
- StaticRunOps(default_ctx, g, state_ptr, 0, idx.num_nodes());
-
- return recording ? state_ptr : OpStatePtr();
-}
-
-
-OpStatePtr CachedOp::DynamicForward(
- const Context& default_ctx,
- const std::vector<NDArray*>& inputs,
- const std::vector<NDArray*>& outputs) {
- using namespace nnvm;
- using namespace imperative;
-
- // Initialize
- bool recording = Imperative::Get()->is_recording();
- auto op_state = OpStatePtr::Create<DynamicRuntime>();
- auto& runtime = op_state.get_state<DynamicRuntime>();
- {
- auto state_ptr = GetCachedOpState(default_ctx);
- auto& state = state_ptr.get_state<CachedOpState>();
- std::lock_guard<std::mutex> lock(state.mutex);
- SetForwardGraph(&state.info, recording, inputs);
- runtime.info.fwd_graph = state.info.fwd_graph;
- }
- nnvm::Graph& g = runtime.info.fwd_graph;
- const auto& idx = g.indexed_graph();
- size_t num_inputs = idx.input_nodes().size();
- auto& buff = runtime.buff;
- auto& states = runtime.op_states;
+ auto op_state_ptr = OpStatePtr::Create<CachedOpState>();
+ auto& cached_op_state = op_state_ptr.get_state<CachedOpState>();
+ auto& buff = cached_op_state.buff;
+ auto& states = cached_op_state.states;
// Allocate entries
states.resize(idx.num_nodes());
@@ -780,98 +446,57 @@ OpStatePtr CachedOp::DynamicForward(
AllocateMemory(g, idx, default_ctx, 0, idx.num_node_entries(),
mem_plan, arrays, &array_reqs);
- const auto& dtypes = g.GetAttr<DTypeVector>("dtype");
- const auto& shapes = g.GetAttr<ShapeVector>("shape");
- const auto& stypes = g.GetAttr<StorageTypeVector>("storage_type");
-
- for (size_t i = 0; i < outputs.size(); ++i) {
- auto eid = idx.entry_id(idx.outputs()[i]);
- arrays[eid] = outputs[i];
- if (!outputs[i]->is_none()) continue;
- *outputs[i] = NDArray(static_cast<NDArrayStorageType>(stypes[eid]),
- shapes[eid], default_ctx, true, dtypes[eid]);
- }
-
const auto& dispatch_modes = g.GetAttr<DispatchModeVector>("dispatch_mode");
if (recording && !inlining_) Imperative::Get()->set_is_recording(false);
+ int prev_bulk_size = Engine::Get()->set_bulk_size(config_.forward_bulk_size);
- RunGraph(false, idx, arrays, 0, idx.num_nodes(), std::move(array_reqs),
- std::move(ref_count), &states, dispatch_modes);
+ Imperative::Get()->RunGraph(
+ false, idx, arrays, 0, idx.num_nodes(), std::move(array_reqs),
+ std::move(ref_count), &states, dispatch_modes);
+ Engine::Get()->set_bulk_size(prev_bulk_size);
Imperative::Get()->set_is_recording(recording);
- return op_state;
-}
-
-void CachedOp::Forward(
- const std::shared_ptr<CachedOp>& op_ptr,
- const std::vector<NDArray*>& inputs,
- const std::vector<NDArray*>& outputs) {
- static const auto cached_op = nnvm::Op::Get("_CachedOp");
-
- CHECK_EQ(inputs.size(), num_inputs());
-
- Context default_ctx = inputs[0]->ctx();
-
- const auto& idx = fwd_graph_.indexed_graph();
- for (size_t i = 0; i < inputs.size(); ++i) {
- CHECK_EQ(inputs[i]->ctx(), default_ctx)
- << "CachedOp 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 " << inputs[i]->ctx();
- }
-
- int prev_bulk_size = Engine::Get()->set_bulk_size(config_.forward_bulk_size);
-
- OpStatePtr op_state;
- if (config_.static_alloc) {
- op_state = StaticForward(default_ctx, inputs, outputs);
- } else {
- op_state = DynamicForward(default_ctx, inputs, outputs);
+ for (size_t i = 0; i < idx.num_node_entries(); ++i) {
+ if (arrays[i] == &buff[i]) continue;
+ buff[i].shape_ = arrays[i]->shape_;
+ buff[i].dtype_ = arrays[i]->dtype_;
+ buff[i].storage_type_ = arrays[i]->storage_type_;
}
- Engine::Get()->set_bulk_size(prev_bulk_size);
-
- if (Imperative::Get()->is_recording() && !inlining_) {
+ if (recording && !inlining_) {
nnvm::NodeAttrs attrs;
attrs.op = cached_op;
attrs.name = "_cachedop";
attrs.parsed = op_ptr;
Imperative::Get()->RecordOp(
- std::move(attrs), inputs, outputs, op_state,
+ std::move(attrs), inputs, outputs, op_state_ptr,
&save_inputs(), &save_outputs());
}
}
-void CachedOp::DynamicBackward(
+void Imperative::CachedOp::Backward(
const bool retain_graph,
- const OpStatePtr& op_state,
+ const OpStatePtr& state,
const std::vector<NDArray*>& inputs,
const std::vector<OpReqType>& reqs,
const std::vector<NDArray*>& outputs) {
using namespace nnvm;
using namespace imperative;
+ CHECK(!Imperative::Get()->is_recording())
+ << "CachedOp does not support higher order gradients. "
+ << "If you want to do backward with create_graph=True please "
+ << "do not use hybridize.";
// Initialize
- Context default_ctx = outputs[0]->ctx();
- auto& runtime = op_state.get_state<DynamicRuntime>();
- {
- auto state_ptr = GetCachedOpState(default_ctx);
- auto& state = state_ptr.get_state<CachedOpState>();
- std::lock_guard<std::mutex> lock(state.mutex);
- state.info.fwd_graph = runtime.info.fwd_graph;
- SetBackwardGraph(&state.info, reqs, inputs);
- runtime.info.full_graph = state.info.full_graph;
- runtime.info.bwd_input_eid = state.info.bwd_input_eid;
- }
- nnvm::Graph& g = runtime.info.full_graph;
+ nnvm::Graph g = GetBackwardGraph(state, reqs, inputs);
const auto& idx = g.indexed_graph();
- auto& buff = runtime.buff;
- auto& states = runtime.op_states;
+
+ auto& cached_op_state = state.get_state<CachedOpState>();
+ auto& buff = cached_op_state.buff;
+ auto& states = cached_op_state.states;
size_t num_forward_outputs = fwd_graph_.outputs.size();
size_t num_forward_nodes = fwd_graph_.indexed_graph().num_nodes();
@@ -881,7 +506,7 @@ void CachedOp::DynamicBackward(
arrays.reserve(buff.size());
for (size_t i = 0; i < buff.size(); ++i) arrays.push_back(&buff[i]);
for (size_t i = 0; i < inputs.size(); ++i) {
- arrays[runtime.info.bwd_input_eid[i]] = inputs[i];
+ arrays[bwd_input_eid_[i]] = inputs[i];
}
for (size_t i = 0, j = num_forward_outputs; i < reqs.size(); ++i) {
if (reqs[i] == kNullOp) continue;
@@ -905,14 +530,20 @@ void CachedOp::DynamicBackward(
if (ref_count[i] == 0) array_reqs[i] = kNullOp;
}
+ Context default_ctx = outputs[0]->ctx();
const auto& mem_plan = g.GetAttr<MemoryPlanVector >("backward_mem_plan");
AllocateMemory(g, idx, default_ctx, num_forward_entries, idx.num_node_entries(),
mem_plan, arrays, &array_reqs);
const auto& dispatch_modes = g.GetAttr<DispatchModeVector>("dispatch_mode");
- RunGraph(retain_graph, idx, arrays, num_forward_nodes, idx.num_nodes(),
- std::move(array_reqs), std::move(ref_count), &states, dispatch_modes);
+ int prev_bulk_size = Engine::Get()->set_bulk_size(config_.backward_bulk_size);
+
+ Imperative::Get()->RunGraph(
+ retain_graph, idx, arrays, num_forward_nodes, idx.num_nodes(),
+ std::move(array_reqs), std::move(ref_count), &states, dispatch_modes);
+
+ Engine::Get()->set_bulk_size(prev_bulk_size);
if (retain_graph) {
buff.resize(num_forward_entries);
@@ -922,99 +553,6 @@ void CachedOp::DynamicBackward(
}
}
-void CachedOp::StaticBackward(
- const bool retain_graph,
- const OpStatePtr& state_ptr,
- const std::vector<NDArray*>& inputs,
- const std::vector<OpReqType>& reqs,
- const std::vector<NDArray*>& outputs) {
- using namespace nnvm;
- using namespace imperative;
-
- Context default_ctx = outputs[0]->ctx();
-
- auto& state = state_ptr.get_state<CachedOpState>();
- std::lock_guard<std::mutex> lock(state.mutex);
-
- bool match = SetBackwardGraph(&state.info, reqs, inputs, true);
-
- nnvm::Graph& g = state.info.full_graph;
- const auto& idx = g.indexed_graph();
- auto num_forward_nodes = state.info.fwd_graph.indexed_graph().num_nodes();
-
- if (!state.bwd_alloc || !match) {
- StaticAllocMemory(state_ptr, true, true);
- }
-
- if (config_.static_shape) {
- for (auto i : config_.param_indices) {
- const auto iter = fwd_input_to_grad_output_.find(i);
- if (iter == fwd_input_to_grad_output_.end()) continue;
- auto entry = grad_graph_.outputs[iter->second];
- if (!idx.exist(entry.node.get())) continue;
- auto eid = idx.entry_id(entry);
- if (!state.arrays[eid]->IsSame(*outputs[iter->second]) ||
- !(state.array_reqs[eid] == reqs[iter->second])) {
- match = false;
- state.array_reqs[eid] = reqs[iter->second];
- *state.arrays[eid] = *outputs[iter->second];
- state.dynamic_entries[eid] = false;
- }
- }
- for (auto i : config_.data_indices) {
- const auto iter = fwd_input_to_grad_output_.find(i);
- if (iter == fwd_input_to_grad_output_.end()) continue;
- auto entry = grad_graph_.outputs[iter->second];
- if (!idx.exist(entry.node.get())) continue;
- auto eid = idx.entry_id(entry);
- state.array_reqs[eid] = reqs[iter->second];
- state.arrays[eid] = outputs[iter->second];
- }
- } else {
- for (size_t i = 0; i < grad_graph_.outputs.size(); ++i) {
- auto entry = grad_graph_.outputs[i];
- if (!idx.exist(entry.node.get())) continue;
- auto eid = idx.entry_id(entry);
- state.array_reqs[eid] = reqs[i];
- state.arrays[eid] = outputs[i];
- }
- }
-
- if (!state.bwd_exec_init || !match) {
- StaticInitExec(state_ptr, true, true);
- }
-
- for (size_t i = 0; i < state.info.bwd_input_eid.size(); ++i) {
- auto eid = state.info.bwd_input_eid[i];
- if (state.dynamic_entries[eid]) state.arrays[eid] = inputs[i];
- }
-
- StaticRunOps(default_ctx, g, state_ptr, num_forward_nodes, idx.num_nodes());
-}
-
-void CachedOp::Backward(
- const bool retain_graph,
- const OpStatePtr& state,
- const std::vector<NDArray*>& inputs,
- const std::vector<OpReqType>& reqs,
- const std::vector<NDArray*>& outputs) {
- using namespace imperative;
- CHECK(!Imperative::Get()->is_recording())
- << "CachedOp does not support higher order gradients. "
- << "If you want to do backward with create_graph=True please "
- << "do not use hybridize.";
-
- int prev_bulk_size = Engine::Get()->set_bulk_size(config_.backward_bulk_size);
-
- if (config_.static_alloc) {
- StaticBackward(retain_graph, state, inputs, reqs, outputs);
- } else {
- DynamicBackward(retain_graph, state, inputs, reqs, outputs);
- }
-
- Engine::Get()->set_bulk_size(prev_bulk_size);
-}
-
NNVM_REGISTER_OP(_CachedOp)
.set_num_inputs([](const NodeAttrs& attrs) {
diff --git a/src/imperative/cached_op.h b/src/imperative/cached_op.h
deleted file mode 100644
index 60a40c5..0000000
--- a/src/imperative/cached_op.h
+++ /dev/null
@@ -1,174 +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.
- */
-
-#ifndef MXNET_IMPERATIVE_CACHED_OP_H_
-#define MXNET_IMPERATIVE_CACHED_OP_H_
-
-#include <mxnet/imperative.h>
-#include <vector>
-#include <atomic>
-#include <utility>
-#include <string>
-#include <unordered_map>
-
-namespace mxnet {
-/*! \brief CachedOp Parameters */
-struct CachedOpConfig : public dmlc::Parameter<CachedOpConfig> {
- uint32_t inline_limit;
- uint32_t forward_bulk_size;
- uint32_t backward_bulk_size;
- bool static_alloc;
- bool static_shape;
- nnvm::Tuple<uint32_t> data_indices;
- nnvm::Tuple<uint32_t> param_indices;
- DMLC_DECLARE_PARAMETER(CachedOpConfig) {
- DMLC_DECLARE_FIELD(static_alloc)
- .set_default(false)
- .describe("Statically allocate memory to improve speed. "
- "Memory usage may increase.");
- DMLC_DECLARE_FIELD(static_shape)
- .set_default(false)
- .describe("Optimize for invariant input shapes between iterations. "
- "Must also set static_alloc to True. "
- "Change of input shapes is still allowed but slower.");
- DMLC_DECLARE_FIELD(inline_limit)
- .set_default(2)
- .describe("Maximum number of operators that can be inlined.");
- DMLC_DECLARE_FIELD(forward_bulk_size)
- .set_default(dmlc::GetEnv("MXNET_EXEC_BULK_EXEC_MAX_NODE_TRAIN", 15))
- .describe("Segment size of bulk execution during forward pass.");
- DMLC_DECLARE_FIELD(backward_bulk_size)
- .set_default(dmlc::GetEnv("MXNET_EXEC_BULK_EXEC_MAX_NODE_TRAIN", 15))
- .describe("Segment size of bulk execution during backward pass.");
- DMLC_DECLARE_FIELD(data_indices)
- .set_default(nnvm::Tuple<uint32_t>())
- .describe("Position of argument variables.");
- DMLC_DECLARE_FIELD(param_indices)
- .set_default(nnvm::Tuple<uint32_t>())
- .describe("Position of parameters.");
- }
-};
-
-class CachedOp {
- public:
- CachedOp(
- const nnvm::Symbol& sym,
- const std::vector<std::pair<std::string, std::string> >& flags);
- ~CachedOp();
- uint32_t num_inputs() {
- return fwd_graph_.indexed_graph().input_nodes().size();
- }
- uint32_t num_outputs() {
- return fwd_graph_.outputs.size();
- }
- uint32_t num_backward_inputs() {
- return bwd_ograd_dep_.size() + bwd_in_dep_.size() + bwd_out_dep_.size();
- }
- std::vector<bool>& save_inputs() {
- return save_inputs_;
- }
- std::vector<bool>& save_outputs() {
- return save_outputs_;
- }
- const std::unordered_set<uint32_t>& mutable_input_nodes() {
- return fwd_graph_.indexed_graph().mutable_input_nodes();
- }
- std::vector<nnvm::NodeEntry> Gradient(
- const nnvm::NodePtr& node,
- const std::vector<nnvm::NodeEntry>& ograds);
- void Forward(
- const std::shared_ptr<CachedOp>& op_ptr,
- const std::vector<NDArray*>& inputs,
- const std::vector<NDArray*>& outputs);
- void Backward(
- const bool retain_graph,
- const OpStatePtr& state,
- const std::vector<NDArray*>& inputs,
- const std::vector<OpReqType>& reqs,
- const std::vector<NDArray*>& outputs);
-
- private:
- struct GraphInfo;
- struct DynamicRuntime;
- struct CachedOpState;
-
- OpStatePtr GetCachedOpState(const Context& ctx);
- bool SetForwardGraph(
- GraphInfo* info,
- const bool recording,
- const std::vector<NDArray*>& inputs);
- bool SetBackwardGraph(
- GraphInfo* info,
- const std::vector<OpReqType>& reqs,
- const std::vector<NDArray*>& inputs,
- bool detect_inplace_addto = false);
- OpStatePtr DynamicForward(
- const Context& default_ctx,
- const std::vector<NDArray*>& inputs,
- const std::vector<NDArray*>& outputs);
- void DynamicBackward(
- const bool retain_graph,
- const OpStatePtr& op_state,
- const std::vector<NDArray*>& inputs,
- const std::vector<OpReqType>& reqs,
- const std::vector<NDArray*>& outputs);
- void StaticAllocMemory(
- const OpStatePtr& state_ptr,
- bool recording,
- bool keep_fwd);
- void StaticInitExec(
- const OpStatePtr& state_ptr,
- bool recording,
- bool keep_fwd);
- void StaticRunOps(
- const Context& default_ctx,
- const nnvm::Graph& g,
- const OpStatePtr& state_ptr,
- size_t start_nid,
- size_t end_nid);
- OpStatePtr StaticForward(
- const Context& default_ctx,
- const std::vector<NDArray*>& inputs,
- const std::vector<NDArray*>& outputs);
- void StaticBackward(
- const bool retain_graph,
- const OpStatePtr& state_ptr,
- const std::vector<NDArray*>& inputs,
- const std::vector<OpReqType>& reqs,
- const std::vector<NDArray*>& outputs);
-
- CachedOpConfig config_;
- nnvm::Graph fwd_graph_;
- nnvm::Graph grad_graph_;
- nnvm::Graph full_graph_;
- bool inlining_;
- std::vector<nnvm::NodeEntry> ograd_entries_;
- std::vector<uint32_t> bwd_in_dep_, bwd_out_dep_, bwd_ograd_dep_;
- std::unordered_map<uint32_t, uint32_t> fwd_input_to_grad_output_;
- std::vector<bool> save_inputs_, save_outputs_;
- std::vector<OpReqType> bwd_output_reqs_;
-
- std::mutex mutex_;
- std::unordered_map<Context, std::vector<OpStatePtr> > cached_op_states_;
-};
-
-using CachedOpPtr = std::shared_ptr<CachedOp>;
-
-} // namespace mxnet
-#endif // MXNET_IMPERATIVE_CACHED_OP_H_
diff --git a/src/imperative/imperative.cc b/src/imperative/imperative.cc
index e165425..7caf305 100644
--- a/src/imperative/imperative.cc
+++ b/src/imperative/imperative.cc
@@ -19,7 +19,6 @@
#include <unordered_set>
#include <iostream>
#include "./imperative_utils.h"
-#include "./cached_op.h"
namespace mxnet {
#if DMLC_CXX11_THREAD_LOCAL
@@ -267,6 +266,95 @@ void Imperative::RecordOp(
}
}
+void Imperative::RunGraph(
+ const bool retain_graph,
+ const nnvm::IndexedGraph& idx,
+ const std::vector<NDArray*> arrays,
+ size_t node_start, size_t node_end,
+ std::vector<OpReqType>&& array_reqs,
+ std::vector<uint32_t>&& ref_count,
+ std::vector<OpStatePtr> *p_states,
+ const DispatchModeVector &dispatch_modes) {
+ using namespace nnvm;
+ using namespace imperative;
+ static auto& createop = nnvm::Op::GetAttr<FCreateOpState>("FCreateOpState");
+ static auto& is_layer_backward = Op::GetAttr<bool>("TIsLayerOpBackward");
+ static const auto bwd_cached_op = Op::Get("_backward_CachedOp");
+
+ std::vector<OpStatePtr>& states = *p_states;
+ bool recording = is_recording();
+
+ std::vector<NDArray*> ndinputs, ndoutputs;
+ ShapeVector arg_shapes;
+ DTypeVector arg_dtypes;
+ std::vector<OpReqType> 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 (node.source->op() == bwd_cached_op) {
+ const auto& cached_op = dmlc::get<CachedOpPtr>(node.source->attrs.parsed);
+ nnvm::Node* fwd_node = node.source->control_deps[0].get();
+ auto fwd_node_id = idx.node_id(fwd_node);
+ cached_op->Backward(retain_graph, states[fwd_node_id], ndinputs, req, ndoutputs);
+ } else 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);
+ InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, req, dispatch_mode, states[i]);
+ if (recording) RecordOp(NodeAttrs(node.source->attrs), ndinputs, ndoutputs, 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);
+ InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs,
+ req, dispatch_mode, states[fwd_node_id]);
+ if (recording) {
+ RecordOp(NodeAttrs(node.source->attrs), ndinputs, ndoutputs, states[fwd_node_id]);
+ }
+ } else {
+ InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, req, dispatch_mode);
+ if (recording) RecordOp(NodeAttrs(node.source->attrs), ndinputs, ndoutputs);
+ }
+
+ for (const auto& j : node.inputs) {
+ size_t eid = idx.entry_id(j);
+ --ref_count[eid];
+ if (ref_count[eid] == 0) arrays[eid]->ptr_.reset();
+ }
+ for (size_t j = 0; j < ndoutputs.size(); ++j) {
+ size_t eid = idx.entry_id(i, j);
+ if (ref_count[eid] == 0) arrays[eid]->ptr_.reset();
+ }
+ }
+}
+
+
std::vector<NDArray*> Imperative::Backward(
const std::vector<NDArray*>& outputs,
const std::vector<NDArray*>& ograds,
diff --git a/src/imperative/imperative_utils.cc b/src/imperative/imperative_utils.cc
deleted file mode 100644
index 464aefc..0000000
--- a/src/imperative/imperative_utils.cc
+++ /dev/null
@@ -1,120 +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.
- */
-
-#include "./imperative_utils.h"
-#include "./cached_op.h"
-
-namespace mxnet {
-namespace imperative {
-void RunGraph(
- const bool retain_graph,
- const nnvm::IndexedGraph& idx,
- const std::vector<NDArray*> arrays,
- size_t node_start, size_t node_end,
- std::vector<OpReqType>&& array_reqs,
- std::vector<uint32_t>&& ref_count,
- std::vector<OpStatePtr> *p_states,
- const DispatchModeVector &dispatch_modes) {
- using namespace nnvm;
- using namespace imperative;
- static auto& createop = nnvm::Op::GetAttr<FCreateOpState>("FCreateOpState");
- static auto& is_layer_backward = Op::GetAttr<bool>("TIsLayerOpBackward");
- static const auto bwd_cached_op = Op::Get("_backward_CachedOp");
-
- const auto imp = Imperative::Get();
-
- std::vector<OpStatePtr>& states = *p_states;
- bool recording = imp->is_recording();
-
- std::vector<NDArray*> ndinputs, ndoutputs;
- ShapeVector arg_shapes;
- DTypeVector arg_dtypes;
- std::vector<OpReqType> 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(array_reqs[eid] == kNullOp || !ndoutputs.back()->is_none());
- }
- const Context& ctx = ndoutputs[0]->ctx();
- const DispatchMode dispatch_mode = dispatch_modes[i];
- if (node.source->op() == bwd_cached_op) {
- const auto& cached_op = dmlc::get<CachedOpPtr>(node.source->attrs.parsed);
- nnvm::Node* fwd_node = node.source->control_deps[0].get();
- auto fwd_node_id = idx.node_id(fwd_node);
- cached_op->Backward(retain_graph, states[fwd_node_id], ndinputs, req, ndoutputs);
- } else 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);
- imp->InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, req, dispatch_mode, states[i]);
- if (recording) {
- imp->RecordOp(NodeAttrs(node.source->attrs), ndinputs, ndoutputs, 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);
- imp->InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs,
- req, dispatch_mode, states[fwd_node_id]);
- if (recording) {
- imp->RecordOp(NodeAttrs(node.source->attrs), ndinputs, ndoutputs, states[fwd_node_id]);
- }
- } else {
- imp->InvokeOp(ctx, node.source->attrs, ndinputs, ndoutputs, req, dispatch_mode);
- if (recording) {
- imp->RecordOp(NodeAttrs(node.source->attrs), ndinputs, ndoutputs);
- }
- }
-
- for (const auto& j : node.inputs) {
- size_t eid = idx.entry_id(j);
- --ref_count[eid];
- if (ref_count[eid] == 0) *arrays[eid] = NDArray();
- }
- for (size_t j = 0; j < ndoutputs.size(); ++j) {
- size_t eid = idx.entry_id(i, j);
- if (ref_count[eid] == 0) *arrays[eid] = NDArray();
- }
- }
-}
-
-} // namespace imperative
-} // namespace mxnet
diff --git a/src/imperative/imperative_utils.h b/src/imperative/imperative_utils.h
index 726531d..06b7e05 100644
--- a/src/imperative/imperative_utils.h
+++ b/src/imperative/imperative_utils.h
@@ -23,7 +23,6 @@
#include <utility>
#include <algorithm>
#include <vector>
-#include <map>
#include <string>
#include "../executor/graph_executor.h"
#include "../executor/exec_pass.h"
@@ -39,24 +38,11 @@ namespace mxnet {
namespace imperative {
struct MemoryPlanInfo {
- int storage_id;
- uint32_t root;
+ uint32_t sid;
size_t size;
bool inplace;
};
-struct EngineOprDeleter {
- void operator()(engine::Opr* handle) {
- Engine::Get()->DeleteOperator(handle);
- }
-};
-
-struct EngineOprSeg {
- bool skip;
- size_t next_nid;
- std::unique_ptr<engine::Opr, EngineOprDeleter> opr;
-};
-
using MemoryPlanVector = std::vector<MemoryPlanInfo>;
inline Context GetContext(const nnvm::NodeAttrs& attrs,
@@ -729,12 +715,10 @@ inline std::vector<Context> PlaceDevice(const nnvm::IndexedGraph& idx) {
inline MemoryPlanVector PlanMemory(
- nnvm::Graph* p_g,
- nnvm::StorageVector&& storage,
+ nnvm::Graph* p_g, nnvm::StorageVector&& storage,
const std::vector<uint32_t>& ref_count,
const std::pair<uint32_t, uint32_t>& node_range = {0, 0},
- const std::pair<uint32_t, uint32_t>& entry_range = {0, 0},
- bool detect_inplace_addto = false) {
+ const std::pair<uint32_t, uint32_t>& entry_range = {0, 0}) {
using namespace nnvm;
nnvm::Graph& g = *p_g;
const auto& idx = g.indexed_graph();
@@ -744,31 +728,31 @@ inline MemoryPlanVector PlanMemory(
g.attrs["ref_count"] = std::make_shared<dmlc::any>(ref_count);
g.attrs["storage"] = std::make_shared<dmlc::any>(std::move(storage));
g = nnvm::ApplyPass(g, "PlanMemory");
- if (detect_inplace_addto) g = exec::DetectInplaceAddTo(g);
const auto& dtypes = g.GetAttr<DTypeVector>("dtype");
const auto& shapes = g.GetAttr<ShapeVector>("shape");
- const auto& storage_inplace = g.GetAttr<std::vector<int> >("storage_inplace_index");
- const auto& storage_ids = g.GetAttr<StorageVector>("storage_id");
+ const auto& stypes = g.GetAttr<StorageTypeVector>("storage_type");
+ auto storage_ids = g.MoveCopyAttr<StorageVector>("storage_id");
+ auto storage_inplace = g.MoveCopyAttr<std::vector<int> >("storage_inplace_index");
uint32_t entry_start = entry_range.first;
uint32_t entry_end =
entry_range.second > entry_start ? entry_range.second : idx.num_node_entries();
MemoryPlanVector mem_plan(idx.num_node_entries());
- std::unordered_map<int, uint32_t> sid_to_root;
+ std::unordered_map<int, uint32_t> sid_to_loc;
for (uint32_t i = entry_start; i < entry_end; ++i) {
+ if (stypes[i] != kDefaultStorage) continue;
if (storage_ids[i] < 0) {
- mem_plan[i] = {storage_ids[i], i, 0, false};
- } else if (!sid_to_root.count(storage_ids[i])) {
+ mem_plan[i] = {i, mshadow::mshadow_sizeof(dtypes[i]) * shapes[i].Size(), false};
+ } else if (!sid_to_loc.count(storage_ids[i])) {
CHECK_LT(storage_inplace[i], 0);
- sid_to_root[storage_ids[i]] = i;
- mem_plan[i] = {storage_ids[i], i,
- mshadow::mshadow_sizeof(dtypes[i]) * shapes[i].Size(),
- false};
+ sid_to_loc[storage_ids[i]] = i;
+ mem_plan[i].sid = i;
+ mem_plan[i].size = mshadow::mshadow_sizeof(dtypes[i]) * shapes[i].Size();
} else {
- uint32_t root = sid_to_root[storage_ids[i]];
- mem_plan[i] = {storage_ids[i], root, 0, storage_inplace[i] >= 0};
- mem_plan[root].size = std::max(mem_plan[root].size,
+ uint32_t loc = sid_to_loc[storage_ids[i]];
+ mem_plan[i] = {loc, 0, storage_inplace[i] >= 0};
+ mem_plan[loc].size = std::max(mem_plan[loc].size,
mshadow::mshadow_sizeof(dtypes[i]) * shapes[i].Size());
}
}
@@ -777,213 +761,39 @@ inline MemoryPlanVector PlanMemory(
}
-inline std::multimap<size_t, NDArray> AllocateMemory(
- const nnvm::Graph& g,
- const nnvm::IndexedGraph& idx,
- const Context& default_ctx,
- const uint32_t entry_start, const uint32_t entry_end,
- const MemoryPlanVector& mem_plan,
- const std::vector<NDArray*>& arrays,
- std::vector<OpReqType> *array_reqs,
- std::multimap<size_t, NDArray>&& pool = std::multimap<size_t, NDArray>()) {
+inline void AllocateMemory(const nnvm::Graph& g,
+ const nnvm::IndexedGraph& idx,
+ const Context& default_ctx,
+ const uint32_t entry_start, const uint32_t entry_end,
+ const MemoryPlanVector& mem_plan,
+ const std::vector<NDArray*>& arrays,
+ std::vector<OpReqType> *array_reqs) {
using namespace nnvm;
const auto& dtypes = g.GetAttr<DTypeVector>("dtype");
const auto& shapes = g.GetAttr<ShapeVector>("shape");
const auto& stypes = g.GetAttr<StorageTypeVector>("storage_type");
- std::multimap<size_t, NDArray> new_pool;
-
for (uint32_t i = entry_start; i < entry_end; ++i) {
- if (mem_plan[i].storage_id == exec::kExternalStorageID) continue;
- CHECK(arrays[i]->is_none());
- if (mem_plan[i].storage_id == exec::kDynamicStorageID) {
- *arrays[i] = NDArray(static_cast<NDArrayStorageType>(stypes[i]),
- shapes[i], default_ctx, true, dtypes[i]);
- continue;
- }
- CHECK_EQ(stypes[i], kDefaultStorage);
- if (mem_plan[i].root == i) {
- CHECK_GT(mem_plan[i].size, 0);
- auto iter = pool.lower_bound(mem_plan[i].size);
- if (iter != pool.end()) {
- *arrays[i] = iter->second.AsArray(shapes[i], dtypes[i]);
- new_pool.insert(*iter);
- pool.erase(iter);
- } else {
+ if (!arrays[i]->is_none()) continue;
+ if (stypes[i] == kDefaultStorage) {
+ if (mem_plan[i].sid == i) {
+ CHECK_GT(mem_plan[i].size, 0);
NDArray buff(TShape({static_cast<nnvm::dim_t>(mem_plan[i].size)}),
default_ctx, true, mshadow::kUint8);
*arrays[i] = buff.AsArray(shapes[i], dtypes[i]);
- new_pool.insert({mem_plan[i].size, buff});
- }
- } else {
- CHECK_GE(mem_plan[mem_plan[i].root].storage_id, 0);
- *arrays[i] = arrays[mem_plan[i].root]->AsArray(shapes[i], dtypes[i]);
- if (mem_plan[i].inplace && array_reqs->at(i) == kWriteTo) {
- array_reqs->at(i) = kWriteInplace;
- }
- }
- }
-
- return new_pool;
-}
-
-inline void SetupOpExec(
- const nnvm::Graph& g,
- size_t nid,
- const std::shared_ptr<exec::OpExecutor>& exec,
- const std::vector<NDArray*> arrays,
- const std::vector<OpReqType> array_reqs) {
- const auto& idx = g.indexed_graph();
- const auto& inode = idx[nid];
- CHECK_EQ(exec->in_array.size(), 0U);
- CHECK_EQ(exec->out_array.size(), 0U);
- for (const auto& e : inode.inputs) {
- CHECK(!arrays[idx.entry_id(e)]->is_none()) << inode.source->attrs.name;
- exec->in_array.push_back(*arrays[idx.entry_id(e)]);
- }
- for (uint32_t index = 0; index < inode.source->num_outputs(); ++index) {
- uint32_t eid = idx.entry_id(nid, index);
- CHECK(!arrays[eid]->is_none()) << inode.source->attrs.name;
- exec->out_array.push_back(*arrays[eid]);
- exec->req.push_back(array_reqs[eid]);
- }
-
- exec->Setup();
-}
-
-inline Engine::OprHandle CreateEngineOp(
- const Context& default_ctx,
- const std::vector<std::shared_ptr<exec::OpExecutor> >& execs) {
- CHECK_GT(execs.size(), 0);
- std::vector<Engine::VarHandle> use_vars, mutate_vars;
-
- for (const auto& exec : execs) {
- CHECK_GT(exec->out_array.size(), 0);
- CHECK(execs.size() == 1 || exec->exec_type() == ExecType::kSync);
-
- // the variables
- for (const auto& nd : exec->in_array) {
- use_vars.push_back(nd.var());
- }
- for (auto& r : exec->op_ctx.requested) {
- mutate_vars.push_back(r.var);
- }
- for (auto& nd : exec->out_array) {
- mutate_vars.push_back(nd.var());
- }
- if (exec->var() != nullptr) {
- mutate_vars.push_back(exec->var());
- }
- }
-
- // dedup vars
- Engine::Get()->DeduplicateVarHandle(&use_vars, &mutate_vars);
- bool is_gpu = default_ctx.dev_mask() == gpu::kDevMask;
- bool is_async = execs.size() > 1 ? false : execs[0]->exec_type() == ExecType::kAsync;
-
- auto exec_fun = [execs, is_async, is_gpu] (
- RunContext ctx, Engine::CallbackOnComplete on_complete) {
- if (is_async) {
- execs[0]->op_ctx.async_on_complete = on_complete;
- }
- for (const auto& exec : execs) exec->Run(ctx, is_gpu);
- // call on complete only if it is async op
- if (!is_async) {
- if (is_gpu) {
- #if MXNET_USE_CUDA
- // Wait GPU kernel to finish.
- ctx.get_stream<gpu>()->Wait();
- #else
- LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
- #endif
- }
- on_complete();
- }
- };
-
- return Engine::Get()->NewOperator(
- exec_fun, use_vars, mutate_vars, FnProperty::kNormal);
-}
-
-inline void CreateEngineOpSeg(
- const nnvm::IndexedGraph& idx,
- const Context default_ctx,
- const size_t start_nid,
- const size_t end_nid,
- const size_t bulk_size,
- const std::unordered_set<uint32_t>& excludes,
- const std::vector<std::shared_ptr<exec::OpExecutor> >& execs,
- const std::vector<int> skip_plus_node,
- std::vector<EngineOprSeg> *opr_segs) {
- size_t seg_start = start_nid;
- std::vector<std::shared_ptr<exec::OpExecutor> > seg_execs;
- for (size_t nid = start_nid; nid < end_nid; ++nid) {
- const auto& node = idx[nid];
- if (node.source->is_variable()) continue;
- if (skip_plus_node.size() && skip_plus_node[nid]) continue;
- auto& exec = execs[nid];
- bool is_async = exec->exec_type() != ExecType::kSync;
- bool valid = exec->out_array.size() > 0;
-
- // Stop at async nodes and invalid node (due to input/output is not allocated)
- bool stop = is_async || !valid || seg_execs.size() >= bulk_size;
- for (size_t i = 0; i < node.inputs.size() && !stop; ++i) {
- if (excludes.count(idx.entry_id(node.inputs[i]))) stop = true;
- }
- auto num_outputs = node.source->num_outputs();
- for (size_t i = 0; i < num_outputs && !stop; ++i) {
- if (excludes.count(idx.entry_id(nid, i))) stop = true;
- }
-
- // Create opr segment for previous nodes.
- if (stop && nid > seg_start) {
- auto& seg = (*opr_segs)[seg_start];
- if (seg_execs.size()) {
- seg = EngineOprSeg{false, nid};
- seg.opr.reset(CreateEngineOp(default_ctx, seg_execs));
} else {
- seg = EngineOprSeg{true, nid, nullptr};
+ *arrays[i] = arrays[mem_plan[i].sid]->AsArray(shapes[i], dtypes[i]);
+ if (mem_plan[i].inplace && array_reqs->at(i) == kWriteTo) {
+ array_reqs->at(i) = kWriteInplace;
+ }
}
- seg_start = nid;
- seg_execs.clear();
- }
-
- seg_execs.push_back(exec);
-
- auto& seg = (*opr_segs)[nid];
- if (is_async) {
- seg = EngineOprSeg{false, nid + 1};
- seg.opr.reset(CreateEngineOp(default_ctx, seg_execs));
- seg_execs.clear();
- seg_start = nid + 1;
- } else if (!valid) {
- seg = EngineOprSeg{false, nid + 1, nullptr};
- seg_execs.clear();
- seg_start = nid + 1;
- }
- }
- // The last segment
- if (end_nid > seg_start) {
- auto& seg = (*opr_segs)[seg_start];
- if (seg_execs.size()) {
- seg = EngineOprSeg{false, end_nid};
- seg.opr.reset(CreateEngineOp(default_ctx, seg_execs));
} else {
- seg = EngineOprSeg{true, end_nid, nullptr};
+ *arrays[i] = NDArray(static_cast<NDArrayStorageType>(stypes[i]),
+ shapes[i], default_ctx, true, dtypes[i]);
}
}
}
-
-void RunGraph(const bool retain_graph,
- const nnvm::IndexedGraph& idx,
- const std::vector<NDArray*> arrays,
- size_t node_start, size_t node_end,
- std::vector<OpReqType>&& array_reqs,
- std::vector<uint32_t>&& ref_count,
- std::vector<OpStatePtr> *p_states,
- const DispatchModeVector &dispatch_modes);
-
} // namespace imperative
} // namespace mxnet
diff --git a/tests/python/unittest/test_gluon.py b/tests/python/unittest/test_gluon.py
index 5701a5d..451fde2 100644
--- a/tests/python/unittest/test_gluon.py
+++ b/tests/python/unittest/test_gluon.py
@@ -22,7 +22,6 @@ from mxnet.test_utils import assert_almost_equal
from mxnet.ndarray.ndarray import _STORAGE_TYPE_STR_TO_ID
from common import setup_module, with_seed, assertRaises, teardown
import numpy as np
-from numpy.testing import assert_array_equal
from nose.tools import raises, assert_raises
from copy import deepcopy
import warnings
@@ -1125,6 +1124,7 @@ def test_hybrid_multi_context():
net.hybridize()
net(mx.nd.zeros((1, 3, 32, 32), ctx=mx.cpu(0))).asnumpy()
+
@with_seed()
def test_zero_grad():
data = mx.nd.random.uniform(shape=(3,3))
@@ -1137,60 +1137,6 @@ def test_zero_grad():
grad = net.collect_params()['test_zero_grad_weight'].grad()
assert_almost_equal(grad.asnumpy(), grad.asnumpy() * 0)
-def check_hybrid_static_memory(**kwargs):
- x = mx.nd.random.uniform(shape=(2, 3, 32, 32))
- x.attach_grad()
-
- net1 = gluon.model_zoo.vision.get_resnet(
- 1, 18, pretrained=True, prefix='net_', ctx=mx.context.current_context())
- net2 = gluon.model_zoo.vision.get_resnet(
- 1, 18, pretrained=True, prefix='net_', ctx=mx.context.current_context())
- net2.hybridize(**kwargs)
- net1(x)
- net2(x)
-
- def test(net, x):
- with mx.autograd.record():
- y = net(x) + net(x)
- y.backward()
-
- grads = {k: v.grad() for k, v in net.collect_params().items() if v.grad_req != 'null'}
-
- return y, grads
-
- y1, grads1 = test(net1, x)
- y2, grads2 = test(net2, x)
-
- assert_almost_equal(y1.asnumpy(), y2.asnumpy(), rtol=1e-3, atol=1e-5)
- for key in grads1:
- assert_almost_equal(grads1[key].asnumpy(), grads2[key].asnumpy(), rtol=1e-3, atol=1e-5)
-
-def test_hybrid_static_memory():
- check_hybrid_static_memory()
- check_hybrid_static_memory(static_alloc=True)
- check_hybrid_static_memory(static_alloc=True, static_shape=True)
-
-def check_hybrid_static_memory_switching(**kwargs):
- net = gluon.model_zoo.vision.get_resnet(
- 1, 18, pretrained=True, ctx=mx.context.current_context())
- net.hybridize(**kwargs)
-
- x = mx.nd.random.uniform(shape=(4, 3, 32, 32))
- net(x)
- with mx.autograd.record():
- y = net(x)
- y.backward()
- x = mx.nd.random.uniform(shape=(2, 3, 32, 32))
- net(x)
- with mx.autograd.record():
- y = net(x)
- y.backward()
- mx.nd.waitall()
-
-def test_hybrid_static_memory_switching():
- check_hybrid_static_memory_switching()
- check_hybrid_static_memory_switching(static_alloc=True)
- check_hybrid_static_memory_switching(static_alloc=True, static_shape=True)
@with_seed()
def test_hook():
@@ -1285,17 +1231,6 @@ def test_legacy_save_params():
model.load_params('test.params', ctx=mx.cpu())
-def test_hybrid_static_memory_recording():
- net = gluon.model_zoo.vision.get_resnet(
- 1, 18, pretrained=True, ctx=mx.context.current_context())
- net.hybridize(static_alloc=True)
-
- x = mx.nd.random.uniform(shape=(1, 3, 32, 32))
- with mx.autograd.record(True):
- net(x)
- net(x)
-
-
if __name__ == '__main__':
import nose
nose.runmodule()
--
To stop receiving notification emails like this one, please contact
marcoabreu@apache.org.