You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2018/04/10 19:47:14 UTC
[GitHub] piiswrong closed pull request #10374: [MXNET-93] Sparse support for
Custom Op
piiswrong closed pull request #10374: [MXNET-93] Sparse support for Custom Op
URL: https://github.com/apache/incubator-mxnet/pull/10374
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/example/numpy-ops/custom_sparse_sqr.py b/example/numpy-ops/custom_sparse_sqr.py
new file mode 100644
index 00000000000..d39e37eae66
--- /dev/null
+++ b/example/numpy-ops/custom_sparse_sqr.py
@@ -0,0 +1,120 @@
+# 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.
+
+# pylint: skip-file
+import mxnet as mx
+import numpy as np
+
+class Sqr(mx.operator.CustomOp):
+ '''Example of how to use custom op with sparse ndarrays
+ '''
+ def forward(self, is_train, req, in_data, out_data, aux):
+ inp = in_data[0]
+ if inp.stype == 'csr':
+ csr_m = inp.data
+ out = mx.nd.sparse.csr_matrix((csr_m * csr_m, inp.indices, inp.indptr), shape=inp.shape)
+ else:
+ out = inp * inp
+ self.assign(out_data[0], req[0], out)
+
+ def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
+ self.assign(in_grad[0], req[0], 2 * mx.nd.sparse.elemwise_mul(in_data[0], out_grad[0]))
+
+@mx.operator.register("sqr")
+class SqrProp(mx.operator.CustomOpProp):
+ def __init__(self):
+ super(SqrProp, self).__init__(need_top_grad=True)
+
+ def list_arguments(self):
+ return ['data']
+
+ def list_outputs(self):
+ return ['output']
+
+ def infer_shape(self, in_shape):
+ return in_shape, [in_shape[0]], []
+
+ def infer_type(self, in_type):
+ return in_type, [in_type[0]], []
+
+ def infer_storage_type(self, in_stype):
+ '''Infer storage type logic for the forward pass
+ Takes a list of storage types for inputs
+ Returns three lists lists, one for input storage types inferred,
+ second for output storage types inferred and third for aux storage
+ types inferred
+ The in_stype is the list containing storage type for inputs
+ If the input is a dense ndarray then we infer the input
+ and output to be dense. If input is csr then input and output
+ are inferred as csr.
+ '''
+ if in_stype[0] == 'default':
+ return ['default'], ['default'], []
+ return ['csr'], ['csr'], []
+
+ def infer_storage_type_backward(self, ograd_stype, in_stype, out_stype, igrad_stype, aux_stype):
+ '''Infer storage type logic for the backward pass
+ Takes storage type of output gradients(ograd_stype), inputs(in_stype),
+ outputs(out_stype) and aux(aux_stype).
+ Returns inferred storage types in the following order:
+ ograd_stype, in_stype, out_stype, igrad_stype (Storage type for input gradients)
+ and aux_stype.
+ '''
+ if in_stype[0] == 'default':
+ return ['default'], ['default'], ['default'], ['default'], []
+ return ['csr'], ['csr'], ['csr'], ['default'], []
+
+ def create_operator(self, ctx, shapes, dtypes):
+ return Sqr()
+
+#Default
+x = mx.nd.array(np.random.uniform(1, 10, size=(4,10)))
+x.attach_grad(stype='default')
+z = mx.nd.ones((4,10))
+with mx.autograd.record():
+ y = mx.nd.Custom(x, op_type='sqr')
+ y.backward(out_grad=z)
+print("Original ndarray")
+print("--------------")
+print(x.asnumpy())
+print("Squared ndarray")
+print("--------------")
+print(y.asnumpy())
+print("stype of input is {}".format(x.stype))
+print("stype of output is {}".format(y.stype))
+print("Grad ndarray")
+print(x.grad.asnumpy())
+
+#Sparse
+x = mx.nd.array(np.random.uniform(1, 10, size=(4,10)))
+x = x.tostype('csr')
+x.attach_grad(stype='default')
+z = mx.nd.ones((4,10))
+z = z.tostype('csr')
+with mx.autograd.record():
+ y = mx.nd.Custom(x, op_type='sqr')
+ y.backward(out_grad=z)
+print("Original ndarray")
+print("--------------")
+print(x.asnumpy())
+print("Squared ndarray")
+print("--------------")
+print(y.asnumpy())
+print("stype of input is {}".format(x.stype))
+print("stype of output is {}".format(y.stype))
+print("Grad ndarray")
+print(x.grad.asnumpy())
diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h
index 9e3f2511278..929ed70bdcc 100644
--- a/include/mxnet/c_api.h
+++ b/include/mxnet/c_api.h
@@ -151,7 +151,9 @@ enum CustomOpPropCallbacks {
kCustomOpPropInferShape,
kCustomOpPropDeclareBackwardDependency,
kCustomOpPropCreateOperator,
- kCustomOpPropInferType
+ kCustomOpPropInferType,
+ kCustomOpPropInferStorageType,
+ kCustomOpPropBackwardInferStorageType
};
@@ -162,6 +164,11 @@ typedef int (*CustomOpDelFunc)(void* /*state*/);
typedef int (*CustomOpListFunc)(char*** /*args*/, void* /*state*/);
typedef int (*CustomOpInferShapeFunc)(int /*num_input*/, int* /*ndims*/,
unsigned** /*shapes*/, void* /*state*/);
+typedef int (*CustomOpInferStorageTypeFunc)(int /*num_input*/, int* /*stypes*/, void* /*state*/);
+typedef int (*CustomOpBackwardInferStorageTypeFunc)(int /*num_input*/,
+ int * /*stypes*/,
+ int * /*tags*/,
+ void * /*state*/);
typedef int (*CustomOpInferTypeFunc)(int /*num_input*/, int* /*types*/, void* /*state*/);
typedef int (*CustomOpBwdDepFunc)(const int* /*out_grad*/, const int* /*in_data*/,
const int* /*out_data*/, int* /*num_deps*/,
diff --git a/include/mxnet/ndarray.h b/include/mxnet/ndarray.h
index b8b7f20fcf3..bec4c0fe303 100644
--- a/include/mxnet/ndarray.h
+++ b/include/mxnet/ndarray.h
@@ -507,6 +507,43 @@ class NDArray {
ret.reuse_ = true;
return ret;
}
+
+ /*!
+ * \brief Update ndarray chunk storage handles using existing ndarray storage handles
+ * Also update the aux_handle, aux_shapes and aux_types.
+ * This is specifically used for custom op to update the inputs and outputs from
+ * the temporary ndarray which stores intermediate custom op results.
+ * Should be used with caution elsewhere. Supports only CSR and RSP formats.
+ */
+ inline void SparseUpdateChunk(const NDArray &arr) const {
+ CHECK(shape_ == arr.shape_) << "ndarray shape is different from the target";
+ CHECK(dtype_ == arr.dtype_) << "ndarray dtype is different from the target";
+ auto stype = arr.storage_type();
+ CHECK(stype == kCSRStorage || stype == kRowSparseStorage)
+ << "Only to be used with CSR and RSP storage types";
+ // swap shandles between src and dst
+ Storage::Handle shandle_dst = arr.ptr_->shandle;
+ arr.ptr_->shandle = ptr_->shandle;
+ ptr_->shandle = shandle_dst;
+
+ ptr_->storage_shape = arr.ptr_->storage_shape;
+ ptr_->storage_type = arr.ptr_->storage_type;
+ ptr_->ctx = arr.ptr_->ctx;
+
+ // swap aux_handles between src and dst
+ size_t aux_idx = 0;
+ CHECK(ptr_->aux_handles.size() == arr.ptr_->aux_handles.size())
+ << "ndarray number of aux_handles is different from target";
+ for (auto &aux_handle : arr.ptr_->aux_handles) {
+ Storage::Handle aux_dst = ptr_->aux_handles[aux_idx];
+ ptr_->aux_handles[aux_idx] = aux_handle;
+ aux_handle = aux_dst;
+ aux_idx++;
+ }
+ ptr_->aux_types = arr.ptr_->aux_types;
+ ptr_->aux_shapes = arr.ptr_->aux_shapes;
+ }
+
/*!
* \brief Get an reshaped NDArray
* \param shape new shape
diff --git a/python/mxnet/operator.py b/python/mxnet/operator.py
index 9bff1176590..60895bdca7b 100644
--- a/python/mxnet/operator.py
+++ b/python/mxnet/operator.py
@@ -16,7 +16,7 @@
# under the License.
# coding: utf-8
-# pylint: disable=invalid-name, protected-access, too-many-arguments, no-self-use, too-many-locals, broad-except
+# pylint: disable=invalid-name, protected-access, too-many-arguments, no-self-use, too-many-locals, broad-except, too-many-lines
"""numpy interface for operators."""
from __future__ import absolute_import
@@ -32,6 +32,10 @@
from .base import c_str, mx_uint, mx_float, ctypes2numpy_shared, NDArrayHandle, py_str
from . import symbol, context
from .ndarray import NDArray, _DTYPE_NP_TO_MX, _DTYPE_MX_TO_NP
+from .ndarray.ndarray import _STORAGE_TYPE_STR_TO_ID, _STORAGE_TYPE_ID_TO_STR
+from .ndarray.ndarray import _STORAGE_TYPE_UNDEFINED, _STORAGE_TYPE_DEFAULT
+from .ndarray.ndarray import _STORAGE_TYPE_CSR, _STORAGE_TYPE_ROW_SPARSE
+from .ndarray import _ndarray_cls
c_int_p = POINTER(c_int)
@@ -522,6 +526,90 @@ def infer_type(self, in_type):
return in_type, [in_type[0]]*len(self.list_outputs()), \
[in_type[0]]*len(self.list_auxiliary_states())
+ def infer_storage_type(self, in_stype):
+ """infer_storage_type interface. Used to infer storage type of
+ inputs and outputs in the forward pass. When this interface is not implemented,
+ all stypes will be inferred as default.
+
+ Parameters
+ ----------
+ in_stype : list of stypes, valid stypes are default, row_sparse and
+ csr
+
+ Returns
+ -------
+ in_stype : list
+ list of argument stypes.
+ out_stype : list
+ list of output types calculated from in_stype,
+ in the same order as declared in list_outputs.
+ aux_type : Optional, list
+ list of aux types calculated from in_stype,
+ in the same order as declared in list_auxiliary_states.
+ """
+ for i, stype in enumerate(in_stype):
+ assert stype == _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_DEFAULT], \
+ "Default infer_storage_type implementation doesnt allow non default stypes: " \
+ "found non default stype '%s' for in_stype[%d]. Please implement " \
+ "infer_storage_type and infer_storage_type_backward interface " \
+ "in your custom operator if you have non-default input stypes" % (stype, i)
+ return in_stype, [in_stype[0]]*len(self.list_outputs()), \
+ [in_stype[0]]*len(self.list_auxiliary_states())
+
+ def infer_storage_type_backward(self, ograd_stype, in_stype, out_stype, igrad_stype, aux_stype):
+ """infer_storage_type_backward interface. Used to infer storage
+ type of inputs and outputs in the backward pass.
+
+ Will raise an error if undefined storage type is returned.
+ Returned lists have to be the same size as the input lists to infer_storage_type_backward,
+ otherwise an exception will be thrown. When this interface is not implemented,
+ all stypes will be inferred as default.
+
+ Parameters
+ ----------
+ ograd_stype : list
+ list of output gradient storage types
+ in_stype : list
+ list of input storage types
+ out_stype : list
+ list of output storage types
+ igrad_stype : list
+ list of input gradient storage types
+ aux_stype : list
+ list of auxiliary storage types
+
+ Returns
+ -------
+ ograd_stype : list
+ list of inferred output gradient storage types
+ in_stype : list
+ list of inferred input storage types
+ out_stype : list
+ list of inferred output storage types
+ igrad_stype : list
+ list of inferred input gradient storage types
+ aux_stype : list
+ list of inferred storage types for auxiliary states
+ """
+ for i, stype in enumerate(ograd_stype):
+ assert stype == _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_DEFAULT], \
+ "Default infer_storage_type_backward implementation doesnt allow non default stypes: " \
+ "found non default stype '%s' for ograd_stype[%d]. Please implement " \
+ "infer_storage_type and infer_storage_type_backward interface " \
+ "in your custom operator if you have non-default output gradient stypes" % (stype, i)
+ for i, stype in enumerate(igrad_stype):
+ if stype == _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_UNDEFINED]:
+ stype = _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_DEFAULT]
+ assert stype == _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_DEFAULT], \
+ "Default infer_storage_type_backward implementation doesnt allow non default stypes: " \
+ "found non default stype '%s' for igrad_stype[%d]. Please implement " \
+ "infer_storage_type and infer_storage_type_backward interface " \
+ "in your custom operator if you have non-default input gradient stypes" % (stype, i)
+ stype_lists = [ograd_stype, in_stype, out_stype, igrad_stype, aux_stype]
+ for stype_list in stype_lists:
+ stype_list[:] = len(stype_list) * [_STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_DEFAULT]]
+ return stype_lists[0], stype_lists[1], stype_lists[2], stype_lists[3], stype_lists[4]
+
def list_outputs(self):
"""list_outputs interface. Can override when creating new operators.
@@ -587,6 +675,7 @@ class _Registry(object):
def __init__(self):
self.ref_holder = {}
self.counter = 0
+ self.result_deps = set()
self.lock = Lock()
def inc(self):
@@ -610,6 +699,9 @@ def do_register(prop_cls):
infershape_functype = CFUNCTYPE(c_int, c_int, POINTER(c_int),
POINTER(POINTER(mx_uint)), c_void_p)
infertype_functype = CFUNCTYPE(c_int, c_int, POINTER(c_int), c_void_p)
+ inferstorage_functype = CFUNCTYPE(c_int, c_int, POINTER(c_int), c_void_p)
+ inferstorage_backward_functype = CFUNCTYPE(c_int, c_int, POINTER(c_int), \
+ POINTER(c_int), c_void_p)
list_functype = CFUNCTYPE(c_int, POINTER(POINTER(POINTER(c_char))), c_void_p)
deps_functype = CFUNCTYPE(c_int, c_int_p, c_int_p, c_int_p,
c_int_p, POINTER(c_int_p), c_void_p)
@@ -665,6 +757,104 @@ def infer_shape_entry(num_tensor, tensor_dims,
return False
return True
+
+ def infer_storage_type_backward_entry(num_tensor, tensor_stypes, tags, _):
+ # pylint: disable=C0301
+ """C Callback for CustomOpProp::InferStorageTypeBackward"""
+ try:
+ tensors = [[] for i in range(5)]
+ for i in range(num_tensor):
+ tensors[tags[i]].append(_STORAGE_TYPE_ID_TO_STR[tensor_stypes[i]])
+ # Ordering of stypes: ograd, input, output, igrad, aux
+ tensors = [tensors[3], tensors[0], tensors[1], tensors[2], tensors[4]]
+ ret = op_prop.infer_storage_type_backward(tensors[0],
+ tensors[1],
+ tensors[2],
+ tensors[3],
+ tensors[4])
+ if len(ret) == 4:
+ ret += []
+ elif len(ret) == 5:
+ pass
+ else:
+ raise AssertionError("infer_storage_type_backward must return 4 or 5 lists")
+ assert len(ret[0]) == len(tensors[0]), \
+ "InferStorageTypeBackward Error: expecting == %d " \
+ "entries in returned output gradient " \
+ "stypes, got %d."%(len(tensors[0]), len(ret[0]))
+ assert len(ret[1]) == len(tensors[1]), \
+ "InferStorageTypeBackward Error: expecting == %d " \
+ "entries in returned input stypes, " \
+ "got %d."%(len(tensors[1]), len(ret[1]))
+ assert len(ret[2]) == len(tensors[2]), \
+ "InferStorageTypeBackward Error: expecting == %d " \
+ "entries in returned output stypes, " \
+ "got %d."%(len(tensors[2]), len(ret[2]))
+ assert len(ret[3]) == len(tensors[3]), \
+ "InferStorageTypeBackward Error: expecting == %d " \
+ "entries in returned input gradient stypes, " \
+ "got %d."%(len(tensors[3]), len(ret[3]))
+ assert len(ret[4]) == len(tensors[4]), \
+ "InferStorageTypeBackward Error: expecting == %d " \
+ "entries in returned aux stypes, " \
+ "got %d."%(len(tensors[4]), len(ret[4]))
+ rstype = []
+ for i, ret_list in enumerate(ret):
+ rstype.extend(ret_list)
+
+ for i, stype in enumerate(rstype):
+ assert stype != _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_UNDEFINED], \
+ "stype should not be undefined"
+ assert stype in _STORAGE_TYPE_STR_TO_ID, \
+ "Provided stype: %s is not valid " \
+ "valid stypes are %s, %s, %s"%(stype,
+ _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_DEFAULT],
+ _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_ROW_SPARSE],
+ _STORAGE_TYPE_ID_TO_STR[_STORAGE_TYPE_CSR])
+ tensor_stypes[i] = _STORAGE_TYPE_STR_TO_ID[stype]
+
+ infer_storage_type_backward_entry._ref_holder = [tensor_stypes]
+ except Exception:
+ print('Error in %s.infer_type: %s' % (reg_name, traceback.format_exc()))
+ return False
+ return True
+
+ def infer_storage_type_entry(num_tensor, tensor_stypes, _):
+ """C Callback for CustomOpProp::InferStorageType"""
+ try:
+ n_in = len(op_prop.list_arguments())
+ n_out = len(op_prop.list_outputs())
+ n_aux = len(op_prop.list_auxiliary_states())
+ assert num_tensor == n_in + n_out + n_aux
+
+ stypes = [_STORAGE_TYPE_ID_TO_STR[tensor_stypes[i]] for i in range(n_in)]
+ ret = op_prop.infer_storage_type(stypes)
+ if len(ret) == 2:
+ istype, ostype = ret
+ astype = []
+ elif len(ret) == 3:
+ istype, ostype, astype = ret
+ else:
+ raise AssertionError("infer_storage_type must return 2 or 3 lists")
+
+ assert len(ostype) == n_out, \
+ "InferStorageType Error: expecting %d entries in returned output " \
+ "stypes, got %d."%(n_out, len(ostype))
+ assert len(istype) == n_in, \
+ "InferStorageType Error: expecting %d entries in returned input " \
+ "stypes, got %d."%(n_in, len(istype))
+ assert len(astype) == n_aux, \
+ "InferStorageType Error: expecting %d entries in returned aux state " \
+ "stypes, got %d."%(n_aux, len(astype))
+ rtype = list(istype) + list(ostype) + list(astype)
+ for i, dtype in enumerate(rtype):
+ tensor_stypes[i] = _STORAGE_TYPE_STR_TO_ID[dtype]
+ infer_storage_type_entry._ref_holder = [tensor_stypes]
+ except Exception:
+ print('Error in %s.infer_type: %s' % (reg_name, traceback.format_exc()))
+ return False
+ return True
+
def infer_type_entry(num_tensor, tensor_types, _):
"""C Callback for CustomOpProp::InferType"""
try:
@@ -684,13 +874,13 @@ def infer_type_entry(num_tensor, tensor_types, _):
raise AssertionError("infer_type must return 2 or 3 lists")
assert len(otype) == n_out, \
"InferType Error: expecting %d entries in returned output " \
- "shapes, got %d."%(n_out, len(otype))
+ "types, got %d."%(n_out, len(otype))
assert len(itype) == n_in, \
"InferType Error: expecting %d entries in returned input " \
- "shapes, got %d."%(n_in, len(itype))
+ "types, got %d."%(n_in, len(itype))
assert len(atype) == n_aux, \
"InferType Error: expecting %d entries in returned aux state " \
- "shapes, got %d."%(n_aux, len(atype))
+ "types, got %d."%(n_aux, len(atype))
rtype = list(itype) + list(otype) + list(atype)
for i, dtype in enumerate(rtype):
tensor_types[i] = _DTYPE_NP_TO_MX[dtype]
@@ -752,6 +942,9 @@ def declare_backward_dependency_entry(out_grad, in_data, out_data, num_dep, deps
out_data = [out_data[i] for i in range(len(op_prop.list_outputs()))]
rdeps = op_prop.declare_backward_dependency(out_grad, in_data, out_data)
num_dep[0] = len(rdeps)
+ _registry.result_deps = set()
+ for dep in rdeps:
+ _registry.result_deps.add(dep)
rdeps = cast(c_array_buf(c_int, array('i', rdeps)), c_int_p)
deps[0] = rdeps
@@ -779,13 +972,13 @@ def forward_entry(num_ndarray, ndarraies, tags, reqs, is_train, _):
tensors = [[] for i in range(5)]
for i in range(num_ndarray):
if tags[i] == 1 or tags[i] == 4:
- tensors[tags[i]].append(NDArray(cast(ndarraies[i],
- NDArrayHandle),
- writable=True))
+ tensors[tags[i]].append(_ndarray_cls(cast(ndarraies[i],
+ NDArrayHandle),
+ writable=True))
else:
- tensors[tags[i]].append(NDArray(cast(ndarraies[i],
- NDArrayHandle),
- writable=False))
+ tensors[tags[i]].append(_ndarray_cls(cast(ndarraies[i],
+ NDArrayHandle),
+ writable=False))
reqs = [req_enum[reqs[i]] for i in range(len(tensors[1]))]
with ctx:
op.forward(is_train=is_train, req=reqs,
@@ -801,15 +994,29 @@ def backward_entry(num_ndarray, ndarraies, tags, reqs, is_train, _):
# pylint: disable=W0613
try:
tensors = [[] for i in range(5)]
+ num_outputs = len(op_prop.list_outputs())
+ num_args = len(op_prop.list_arguments())
for i in range(num_ndarray):
+ if i in _registry.result_deps or i >= (num_outputs * 2 + num_args):
+ # If it is a backward dependency or output or aux:
+ # Set stype as undefined so that it returns
+ # ndarray based on existing stype
+ stype = _STORAGE_TYPE_UNDEFINED
+ else:
+ # If it is some input, output or out grad ndarray not part of
+ # backward dependency it is empty and thus the ndarray should
+ # be set to default
+ stype = _STORAGE_TYPE_DEFAULT
if tags[i] == 2 or tags[i] == 4:
- tensors[tags[i]].append(NDArray(cast(ndarraies[i],
- NDArrayHandle),
- writable=True))
+ tensors[tags[i]].append(_ndarray_cls(cast(ndarraies[i],
+ NDArrayHandle),
+ writable=True,
+ stype=stype))
else:
- tensors[tags[i]].append(NDArray(cast(ndarraies[i],
- NDArrayHandle),
- writable=False))
+ tensors[tags[i]].append(_ndarray_cls(cast(ndarraies[i],
+ NDArrayHandle),
+ writable=False,
+ stype=stype))
reqs = [req_enum[reqs[i]] for i in range(len(tensors[2]))]
with ctx:
op.backward(req=reqs,
@@ -867,7 +1074,9 @@ def delete_entry(_):
infershape_functype(infer_shape_entry),
deps_functype(declare_backward_dependency_entry),
createop_functype(create_operator_entry),
- infertype_functype(infer_type_entry)]
+ infertype_functype(infer_type_entry),
+ inferstorage_functype(infer_storage_type_entry),
+ inferstorage_backward_functype(infer_storage_type_backward_entry)]
callbacks = [cast(i, CFUNCTYPE(c_int)) for i in callbacks]
contexts = [None]*len(callbacks)
ret[0] = MXCallbackList(c_int(len(callbacks)),
diff --git a/src/c_api/c_api_function.cc b/src/c_api/c_api_function.cc
index e8ca18944d6..cea8c9553cc 100644
--- a/src/c_api/c_api_function.cc
+++ b/src/c_api/c_api_function.cc
@@ -72,33 +72,38 @@ OpStatePtr CreateState(const nnvm::NodeAttrs& attrs,
void Forward(const OpStatePtr& state,
const OpContext& ctx,
- const std::vector<TBlob>& inputs,
+ const std::vector<NDArray>& inputs,
const std::vector<OpReqType>& req,
- const std::vector<TBlob>& outputs) {
+ const std::vector<NDArray>& outputs) {
LOG(FATAL) << "Not reached";
}
void Backward(const OpStatePtr& state,
const OpContext& ctx,
- const std::vector<TBlob>& inputs,
+ const std::vector<NDArray>& inputs,
const std::vector<OpReqType>& req,
- const std::vector<TBlob>& outputs) {
+ const std::vector<NDArray>& outputs) {
const CustomFunctionParam& params = state.get_state<CustomFunctionParam>();
std::vector<NDArrayHandle> ptrs;
std::vector<NDArray> cpys;
+ std::vector<int> tags;
+ std::unordered_set<int> input_tags({0});
+ std::unordered_set<int> output_tags({1});
auto dev_id = ctx.run_ctx.ctx.dev_id;
for (const auto& i : inputs) {
- NDArray* nd = new NDArray(i, dev_id);
+ NDArray* nd = new NDArray(i.data(), dev_id);
ptrs.push_back(reinterpret_cast<NDArrayHandle>(nd));
cpys.push_back(*nd);
+ tags.push_back(0);
}
for (const auto& i : outputs) {
- NDArray* nd = new NDArray(i, dev_id);
+ NDArray* nd = new NDArray(i.data(), dev_id);
ptrs.push_back(reinterpret_cast<NDArrayHandle>(nd));
cpys.push_back(*nd);
+ tags.push_back(1);
}
op::custom::CustomOperator::Get()->Push(
@@ -109,9 +114,23 @@ void Backward(const OpStatePtr& state,
const_cast<NDArrayHandle*>(ptrs.data()),
reinterpret_cast<const int*>(req.data()), ctx.is_train,
params.info->contexts[kCustomFunctionBackward]));
- }, ctx, false, ctx.is_train, cpys);
+ }, ctx, false, ctx.is_train, cpys, tags, output_tags, outputs);
}
+inline bool InferStorageType(const nnvm::NodeAttrs& attrs, const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int>* iattr, std::vector<int>* oattr) {
+ using namespace op;
+
+ for (size_t i = 0; i < iattr->size(); ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*iattr, i, kDefaultStorage);
+ }
+ for (size_t i = 0; i < oattr->size(); ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*oattr, i, kDefaultStorage);
+ }
+ DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx);
+ return true;
+}
NNVM_REGISTER_OP(_CustomFunction)
.set_num_inputs([](const NodeAttrs& attrs) {
@@ -138,8 +157,9 @@ NNVM_REGISTER_OP(_CustomFunction)
})
.set_attr<FCreateOpState>("FCreateOpState", CreateState)
.set_attr<nnvm::FGradient>("FGradient", Gradient)
-.set_attr<FStatefulCompute>("FStatefulCompute<cpu>", Forward)
-.set_attr<FStatefulCompute>("FStatefulCompute<gpu>", Forward);
+.set_attr<FStatefulComputeEx>("FStatefulComputeEx<cpu>", Forward)
+.set_attr<FStatefulComputeEx>("FStatefulComputeEx<gpu>", Forward)
+.set_attr<FInferStorageType>("FInferStorageType", InferStorageType);
NNVM_REGISTER_OP(_backward_CustomFunction)
@@ -156,8 +176,9 @@ NNVM_REGISTER_OP(_backward_CustomFunction)
.set_attr<FExecType>("FExecType", [](const NodeAttrs& attrs) {
return ExecType::kAsync;
})
-.set_attr<FStatefulCompute>("FStatefulCompute<cpu>", Backward)
-.set_attr<FStatefulCompute>("FStatefulCompute<gpu>", Backward);
+.set_attr<FStatefulComputeEx>("FStatefulComputeEx<cpu>", Backward)
+.set_attr<FStatefulComputeEx>("FStatefulComputeEx<gpu>", Backward)
+.set_attr<FInferStorageType>("FInferStorageType", InferStorageType);
} // namespace custom_function
} // namespace mxnet
diff --git a/src/imperative/imperative_utils.h b/src/imperative/imperative_utils.h
index 61e246dcbe4..0d6525dce36 100644
--- a/src/imperative/imperative_utils.h
+++ b/src/imperative/imperative_utils.h
@@ -437,19 +437,31 @@ inline void PushOperator(const OpStatePtr& state,
std::vector<NDArray> inputs, outputs;
DerefInputOutput(p_inputs, p_outputs, &inputs, &outputs);
- auto fcompute = common::GetFCompute<FStatefulCompute>(op, "FStatefulCompute", ctx);
- auto fcompute_ex = common::GetFCompute<FStatefulComputeEx>(op, "FStatefulComputeEx", ctx);
+ auto fcompute =
+ common::GetFCompute<FStatefulCompute>(op, "FStatefulCompute", ctx);
+ auto fcompute_ex =
+ common::GetFCompute<FStatefulComputeEx>(op, "FStatefulComputeEx", ctx);
if (fcompute_ex != nullptr && dispatch_mode == DispatchMode::kFComputeEx) {
- CHECK(exec_type == ExecType::kSync);
- Engine::Get()->PushSync(
- [=](RunContext rctx) {
- OpContext opctx{is_train, rctx, engine::CallbackOnComplete(), requested};
- fcompute_ex(state, opctx, inputs, req, outputs);
- if (ctx.dev_mask() == gpu::kDevMask) {
- rctx.get_stream<gpu>()->Wait();
- }
- }, ctx, read_vars, write_vars, FnProperty::kNormal,
- 0, op->name.c_str());
+ const auto& run = [=](RunContext rctx,
+ engine::CallbackOnComplete on_complete) {
+ OpContext opctx{is_train, rctx, on_complete, requested};
+ fcompute_ex(state, opctx, inputs, req, outputs);
+ if (ctx.dev_mask() == gpu::kDevMask && exec_type == ExecType::kSync) {
+ rctx.get_stream<gpu>()->Wait();
+ }
+ };
+
+ if (exec_type == ExecType::kSync) {
+ Engine::Get()->PushSync(
+ [=](RunContext rctx) { run(rctx, engine::CallbackOnComplete()); },
+ ctx, read_vars, write_vars, FnProperty::kNormal, 0,
+ op->name.c_str());
+ } else {
+ CHECK(exec_type == ExecType::kAsync);
+ Engine::Get()->PushAsync(run, ctx, read_vars, write_vars,
+ FnProperty::kAsync, 0,
+ op->name.c_str());
+ }
} else {
CHECK(fcompute != nullptr)
<< "One of FStatefulCompute and FStatefulComputeEx must be registered "
diff --git a/src/operator/custom/custom-inl.h b/src/operator/custom/custom-inl.h
index 13c8d3434b1..de82f3e8eef 100644
--- a/src/operator/custom/custom-inl.h
+++ b/src/operator/custom/custom-inl.h
@@ -64,35 +64,71 @@ class CustomOperator {
return nullptr;
}
- template<typename Func>
- void Push(const Func& func,
- const OpContext& ctx,
- bool recording,
- bool training,
- const std::vector<NDArray>& arrs) {
+ // For sparse the memory allocation is done during execution of operator
+ // which leads to changing of the pointers stored by ndarray chunk.
+ // Thus the changes to the copied ndarries don't propage to final
+ // inputs and outputs unlike the dense case. Passing vector of inputs and
+ // outputs ndarrays as args and updating the inputs and outputs ndarray
+ // chunk pointers to be same as the copied ndarrays.
+ template <typename Func>
+ void Push(const Func& func, const OpContext& ctx, bool recording,
+ bool training, const std::vector<NDArray>& arrs,
+ const std::vector<int>& tags,
+ const std::unordered_set<int>& output_tags,
+ const std::vector<NDArray>& outputs) {
if (naive_engine_) {
func();
+ for (size_t i = 0, out_idx = 0; i < arrs.size(); i++) {
+ if (arrs[i].storage_type() == kDefaultStorage ||
+ arrs[i].storage_type() == kUndefinedStorage)
+ continue;
+ if (output_tags.count(tags[i]) > 0) {
+ outputs[out_idx].SparseUpdateChunk(arrs[i]);
+ out_idx++;
+ }
+ }
ctx.async_on_complete();
return;
}
std::unique_lock<std::mutex> lock(mutex_);
- q_.push(
- [=]() mutable {
- bool prev_recording = Imperative::Get()->set_is_recording(recording);
- bool prev_training = Imperative::Get()->set_is_training(training);
+ q_.push([=]() mutable {
+ bool prev_recording = Imperative::Get()->set_is_recording(recording);
+ bool prev_training = Imperative::Get()->set_is_training(training);
- func();
+ func();
+
+ Imperative::Get()->set_is_training(prev_training);
+ Imperative::Get()->set_is_recording(prev_recording);
- Imperative::Get()->set_is_training(prev_training);
- Imperative::Get()->set_is_recording(prev_recording);
+ std::vector<Engine::VarHandle> vars, vars2;
+ size_t idx = 0;
+ for (const auto& i : arrs) {
+ vars.push_back(i.var());
+ if (output_tags.count(tags[idx]) > 0) {
+ if (i.storage_type() == kDefaultStorage ||
+ i.storage_type() == kUndefinedStorage)
+ continue;
+ vars2.push_back(i.var());
+ idx++;
+ }
+ }
- std::vector<Engine::VarHandle> vars;
- for (const auto& i : arrs) vars.push_back(i.var());
- Engine::Get()->PushSync([=](RunContext rctx) {
+ Engine::Get()->PushSync(
+ [=](RunContext rctx) {
+ for (size_t i = 0, out_idx = 0; i < arrs.size(); i++) {
+ if (arrs[i].storage_type() == kDefaultStorage ||
+ arrs[i].storage_type() == kUndefinedStorage)
+ continue;
+ if (output_tags.count(tags[i]) > 0) {
+ outputs[out_idx].SparseUpdateChunk(arrs[i]);
+ out_idx++;
+ }
+ }
ctx.async_on_complete();
- }, ctx.run_ctx.ctx, vars, {},
- FnProperty::kNormal, 0, "CustomOperator");
- });
+ },
+ ctx.run_ctx.ctx, vars, vars2, FnProperty::kNormal, 0,
+ "CustomOperator");
+ });
cv_.notify_all();
}
diff --git a/src/operator/custom/custom.cc b/src/operator/custom/custom.cc
index 164c2cc597f..d117a284216 100644
--- a/src/operator/custom/custom.cc
+++ b/src/operator/custom/custom.cc
@@ -28,6 +28,7 @@
#include <mxnet/ndarray.h>
#include "../elemwise_op_common.h"
+#include "../operator_common.h"
namespace mxnet {
namespace op {
@@ -45,6 +46,31 @@ struct CustomParam {
std::shared_ptr<MXCallbackList> info;
};
+/*! \brief allocate ndarrays from existing ndarrays
+ */
+inline void AllocateNDArrayCopy(NDArray** nd,
+ const std::vector<NDArray>& inputs,
+ size_t idx, int dev_id) {
+ std::vector<TBlob> aux;
+ NDArrayStorageType stype = inputs[idx].storage_type();
+ switch (stype) {
+ case kUndefinedStorage:
+ case kDefaultStorage:
+ *nd = new NDArray(inputs[idx].data(), dev_id);
+ break;
+ case kRowSparseStorage:
+ aux.push_back(inputs[idx].aux_data(rowsparse::kIdx));
+ *nd = new NDArray(stype, inputs[idx].shape(), inputs[idx].data(), aux,
+ dev_id);
+ break;
+ case kCSRStorage:
+ aux.push_back(inputs[idx].aux_data(csr::kIndPtr));
+ aux.push_back(inputs[idx].aux_data(csr::kIdx));
+ *nd = new NDArray(stype, inputs[idx].shape(), inputs[idx].data(), aux,
+ dev_id);
+ break;
+ }
+}
template<CustomOpPropCallbacks Type>
std::vector<std::string> List(const NodeAttrs& attrs) {
@@ -266,97 +292,249 @@ OpStatePtr CreateState(const NodeAttrs& attrs, Context ctx,
return OpStatePtr::Create<CustomParam>(state);
}
-void Forward(const OpStatePtr& state,
- const OpContext& ctx,
- const std::vector<TBlob>& inputs,
- const std::vector<OpReqType>& req,
- const std::vector<TBlob>& outputs) {
+void ForwardEx(const OpStatePtr& state, const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
const CustomParam& params = state.get_state<CustomParam>();
std::vector<void*> ptrs;
+ // Tags are provided to the callback to provide the frontend
std::vector<int> tags;
std::vector<NDArray> cpys;
+ // info on what ndarray is at each position in the input and output vector
+ // 0 - Input
+ // 1 - Output
+ // 4 - aux
+ std::unordered_set<int> input_tags({0, 4});
+ std::unordered_set<int> output_tags({1});
+
auto dev_id = ctx.run_ctx.ctx.dev_id;
for (size_t i = 0; i < params.num_args; ++i) {
- NDArray *nd = new NDArray(inputs[i], dev_id);
+ NDArray* nd;
+ AllocateNDArrayCopy(&nd, inputs, i, dev_id);
cpys.push_back(*nd);
ptrs.push_back(reinterpret_cast<void*>(nd));
tags.push_back(0);
}
for (size_t i = 0; i < params.num_outs; ++i) {
- NDArray *nd = new NDArray(outputs[i], dev_id);
+ NDArray* nd;
+ AllocateNDArrayCopy(&nd, outputs, i, dev_id);
cpys.push_back(*nd);
ptrs.push_back(reinterpret_cast<void*>(nd));
tags.push_back(1);
}
for (size_t i = 0; i < params.num_auxs; ++i) {
- NDArray *nd = new NDArray(inputs[i+params.num_args], dev_id);
+ size_t idx = i + params.num_args;
+ NDArray* nd;
+ AllocateNDArrayCopy(&nd, inputs, idx, dev_id);
cpys.push_back(*nd);
ptrs.push_back(reinterpret_cast<void*>(nd));
tags.push_back(4);
}
CustomOperator::Get()->Push(
- [=]() {
- CHECK(reinterpret_cast<CustomOpFBFunc>(params.info->callbacks[kCustomOpForward])(
- ptrs.size(), const_cast<void**>(ptrs.data()), const_cast<int*>(tags.data()),
- reinterpret_cast<const int*>(req.data()), static_cast<int>(ctx.is_train),
- params.info->contexts[kCustomOpForward]));
- }, ctx, false, ctx.is_train, cpys);
+ [=]() {
+ CHECK(reinterpret_cast<CustomOpFBFunc>(
+ params.info->callbacks[kCustomOpForward])(
+ ptrs.size(), const_cast<void**>(ptrs.data()),
+ const_cast<int*>(tags.data()),
+ reinterpret_cast<const int*>(req.data()),
+ static_cast<int>(ctx.is_train),
+ params.info->contexts[kCustomOpForward]));
+ },
+ ctx, false, ctx.is_train, cpys, tags, output_tags, outputs);
}
-
-void Backward(const OpStatePtr& state,
- const OpContext& ctx,
- const std::vector<TBlob>& inputs,
- const std::vector<OpReqType>& req,
- const std::vector<TBlob>& outputs) {
+void BackwardEx(const OpStatePtr& state, const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
const CustomParam& params = state.get_state<CustomParam>();
- size_t total = 2*params.num_args + 2*params.num_outs + params.num_auxs;
- std::vector<void*> ptrs(params.num_args + 2*params.num_outs, nullptr);
+ size_t total = 2 * params.num_args + 2 * params.num_outs + params.num_auxs;
+ std::vector<void*> ptrs(params.num_args + 2 * params.num_outs, nullptr);
+
std::vector<int> tags;
std::vector<NDArray> cpys;
ptrs.reserve(total);
tags.reserve(total);
+ cpys.reserve(total);
+
+ // info on what ndarray is at each position in the input and output vector
+ // 3 - out grads
+ // 0 - inputs
+ // 1 - outputs
+ // 4 - auxs
+ // 2 - in grads
+ std::unordered_set<int> input_tags({3, 0, 1, 4});
+ std::unordered_set<int> output_tags({2});
+
for (size_t i = 0; i < params.num_outs; ++i) tags.push_back(3);
for (size_t i = 0; i < params.num_args; ++i) tags.push_back(0);
for (size_t i = 0; i < params.num_outs; ++i) tags.push_back(1);
auto dev_id = ctx.run_ctx.ctx.dev_id;
+
for (size_t i = 0; i < params.bwd_idx.size(); ++i) {
- NDArray *nd = new NDArray(inputs[i], dev_id);
+ NDArray* nd;
+ AllocateNDArrayCopy(&nd, inputs, i, dev_id);
cpys.push_back(*nd);
ptrs[params.bwd_idx[i]] = reinterpret_cast<void*>(nd);
}
for (size_t i = 0; i < ptrs.size(); ++i) {
- if (ptrs[i] == nullptr) ptrs[i] = reinterpret_cast<void*>(new NDArray());
+ NDArray* nd;
+ if (ptrs[i] == nullptr) {
+ nd = new NDArray();
+ ptrs[i] = reinterpret_cast<void*>(nd);
+ }
}
- for (const auto& i : outputs) {
- NDArray* nd = new NDArray(i, dev_id);
+ for (size_t i = 0; i < outputs.size(); ++i) {
+ NDArray* nd;
+ AllocateNDArrayCopy(&nd, outputs, i, dev_id);
cpys.push_back(*nd);
ptrs.push_back(reinterpret_cast<void*>(nd));
tags.push_back(2);
}
+
for (size_t i = 0; i < params.num_auxs; ++i) {
- NDArray* nd = new NDArray(inputs[inputs.size()-params.num_auxs+i], dev_id);
+ size_t idx = inputs.size() - params.num_auxs + i;
+ NDArray* nd;
+ AllocateNDArrayCopy(&nd, inputs, idx, dev_id);
cpys.push_back(*nd);
ptrs.push_back(reinterpret_cast<void*>(nd));
tags.push_back(4);
}
-
CustomOperator::Get()->Push(
[=]() {
CHECK(reinterpret_cast<CustomOpFBFunc>(params.info->callbacks[kCustomOpBackward])(
ptrs.size(), const_cast<void**>(ptrs.data()), const_cast<int*>(tags.data()),
reinterpret_cast<const int*>(req.data()), static_cast<int>(ctx.is_train),
params.info->contexts[kCustomOpBackward]));
- }, ctx, false, ctx.is_train, cpys);
+ }, ctx, false, ctx.is_train, cpys, tags, output_tags, outputs);
+}
+
+// infer storage backward function for custom op which assigns kDefaultStorage for
+// all undefined stypes and dispatches on DispatchMode::kFComputeEx.
+inline bool BackwardInferStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int>* iattr,
+ std::vector<int>* oattr) {
+ const CustomParam& params = nnvm::get<CustomParam>(attrs.parsed);
+
+ if (params.info->num_callbacks <= kCustomOpPropBackwardInferStorageType) {
+ for (size_t i = 0; i < iattr->size(); i++) {
+ STORAGE_TYPE_ASSIGN_CHECK(*iattr, i, kDefaultStorage);
+ }
+ for (size_t i = 0; i < oattr->size(); i++) {
+ STORAGE_TYPE_ASSIGN_CHECK(*oattr, i, kDefaultStorage);
+ }
+ DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx);
+ return true;
+ }
+
+ size_t total = 2 * params.num_args + 2 * params.num_outs + params.num_auxs;
+ size_t bwd_deps_size = params.bwd_idx.size();
+ std::vector<int> stypes(bwd_deps_size, -1);
+ std::vector<int> tags;
+ stypes.reserve(total);
+ tags.reserve(total);
+
+ for (size_t i = 0; i < bwd_deps_size; i++) {
+ if (params.bwd_idx[i] < static_cast<int>(params.num_outs))
+ tags.push_back(3);
+ else if (params.bwd_idx[i] <
+ static_cast<int>(params.num_outs + params.num_args))
+ tags.push_back(0);
+ else
+ tags.push_back(1);
+ stypes[i] = (*iattr)[i];
+ }
+
+ for (size_t i = 0; i < oattr->size(); i++) {
+ stypes.push_back((*oattr)[i]);
+ tags.push_back(2);
+ }
+
+ for (size_t i = (iattr->size() - params.num_auxs); i < iattr->size(); i++) {
+ stypes.push_back((*iattr)[i]);
+ tags.push_back(4);
+ }
+
+ CHECK(reinterpret_cast<CustomOpBackwardInferStorageTypeFunc>(
+ params.info->callbacks[kCustomOpPropBackwardInferStorageType])(
+ stypes.size(), stypes.data(), tags.data(),
+ params.info->contexts[kCustomOpPropBackwardInferStorageType]));
+
+ for (size_t i = 0; i < bwd_deps_size; ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*iattr, i, stypes[i]);
+ }
+ for (size_t i = 0; i < oattr->size(); ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*oattr, i, stypes[i + bwd_deps_size]);
+ }
+ for (size_t i = 0; i < params.num_auxs; ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*iattr, (i + iattr->size() - params.num_auxs),
+ stypes[i + params.num_outs + bwd_deps_size]);
+ }
+
+ DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx);
+ return true;
+}
+
+// infer storage function for custom op which assigns kDefaultStorage for
+// all undefined stypes and dispatches on DispatchMode::kFComputeEx.
+inline bool InferStorageType(const nnvm::NodeAttrs& attrs, const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int>* iattr, std::vector<int>* oattr) {
+ const CustomParam& params = nnvm::get<CustomParam>(attrs.parsed);
+
+ if (params.info->num_callbacks <= kCustomOpPropInferStorageType) {
+ for (size_t i = 0; i < iattr->size(); i++) {
+ STORAGE_TYPE_ASSIGN_CHECK(*iattr, i, kDefaultStorage);
+ }
+ for (size_t i = 0; i < oattr->size(); i++) {
+ STORAGE_TYPE_ASSIGN_CHECK(*oattr, i, kDefaultStorage);
+ }
+ DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx);
+ return true;
+ }
+
+ std::vector<int> stypes;
+ stypes.reserve(params.num_args + params.num_outs + params.num_auxs);
+ for (size_t i = 0; i < params.num_args; ++i) {
+ stypes.push_back((*iattr)[i]);
+ }
+ for (const auto& i : *oattr) {
+ stypes.push_back(i);
+ }
+ for (size_t i = 0; i < params.num_auxs; ++i) {
+ stypes.push_back((*iattr)[params.num_args + i]);
+ }
+
+ CHECK(reinterpret_cast<CustomOpInferStorageTypeFunc>(
+ params.info->callbacks[kCustomOpPropInferStorageType])(
+ stypes.size(), stypes.data(),
+ params.info->contexts[kCustomOpPropInferStorageType]));
+
+ for (size_t i = 0; i < params.num_args; ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*iattr, i, stypes[i]);
+ }
+ for (size_t i = 0; i < params.num_outs; ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*oattr, i, stypes[params.num_args + i]);
+ }
+ for (size_t i = 0; i < params.num_auxs; ++i) {
+ STORAGE_TYPE_ASSIGN_CHECK(*iattr, params.num_args + i,
+ stypes[params.num_args + params.num_outs + i]);
+ }
+
+ DISPATCH_MODE_ASSIGN_CHECK(dispatch_mode, 0, DispatchMode::kFComputeEx);
+ return true;
}
NNVM_REGISTER_OP(Custom)
@@ -396,8 +574,9 @@ Please check the tutorial here: http://mxnet.io/faq/new_op.html.
})
.set_attr<nnvm::FGradient>("FGradient", Gradient)
.set_attr<FCreateOpState>("FCreateOpState", CreateState)
-.set_attr<FStatefulCompute>("FStatefulCompute<cpu>", Forward)
-.set_attr<FStatefulCompute>("FStatefulCompute<gpu>", Forward)
+.set_attr<FStatefulComputeEx>("FStatefulComputeEx<cpu>", ForwardEx)
+.set_attr<FStatefulComputeEx>("FStatefulComputeEx<gpu>", ForwardEx)
+.set_attr<FInferStorageType>("FInferStorageType", InferStorageType)
.add_argument("data", "NDArray-or-Symbol[]", "Input data for the custom operator.")
.add_argument("op_type", "string", "Name of the custom operator. "
"This is the name that is passed to `mx.operator.register` "
@@ -418,8 +597,9 @@ NNVM_REGISTER_OP(_backward_Custom)
.set_attr<FExecType>("FExecType", [](const NodeAttrs& attrs) {
return ExecType::kAsync;
})
-.set_attr<FStatefulCompute>("FStatefulCompute<cpu>", Backward)
-.set_attr<FStatefulCompute>("FStatefulCompute<gpu>", Backward);
+.set_attr<FInferStorageType>("FInferStorageType", BackwardInferStorageType)
+.set_attr<FStatefulComputeEx>("FStatefulComputeEx<cpu>", BackwardEx)
+.set_attr<FStatefulComputeEx>("FStatefulComputeEx<gpu>", BackwardEx);
} // namespace custom
} // namespace op
diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py
index 8d66b2a74ac..87cf5bb699a 100644
--- a/tests/python/unittest/test_operator.py
+++ b/tests/python/unittest/test_operator.py
@@ -4025,12 +4025,22 @@ def test_rcbrt_op():
def test_custom_op():
class Sqr(mx.operator.CustomOp):
def forward(self, is_train, req, in_data, out_data, aux):
- self.assign(out_data[0], req[0], in_data[0]*in_data[0])
- aux[0][:] = 1
+ if in_data[0].stype == 'default':
+ aux[0][:] = 1
+ self.assign(out_data[0], req[0], in_data[0]*in_data[0])
+ else:
+ inp = in_data[0]
+ csr_m = inp.data * inp.data
+ out = mx.nd.sparse.csr_matrix((csr_m, inp.indices, inp.indptr), shape=inp.shape)
+ self.assign(out_data[0], req[0], out)
+ if (in_data[0].stype == 'csr'):
+ assert(isinstance(out_data[0], mx.nd.sparse.CSRNDArray))
+
def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
- self.assign(in_grad[0], req[0], 2*in_data[0]*out_grad[0])
- assert (aux[0].asnumpy() == 1).all()
+ self.assign(in_grad[0], req[0], 2 * mx.nd.sparse.elemwise_mul(in_data[0], out_grad[0]))
+ if in_data[0].stype == 'default':
+ assert (aux[0].asnumpy() == 1).all()
@mx.operator.register("sqr")
class SqrProp(mx.operator.CustomOpProp):
@@ -4052,6 +4062,17 @@ def infer_shape(self, in_shape):
def infer_type(self, in_type):
return in_type, [in_type[0]], [in_type[0]]
+ def infer_storage_type(self, in_stype):
+ if in_stype[0] == 'default':
+ return ['default'], ['default'], ['default']
+ return ['csr'], ['csr'], ['csr']
+
+ def infer_storage_type_backward(self, ograd_stype, in_stype,
+ out_stype, igrad_stype, aux_stype):
+ if in_stype[0] == 'default':
+ return ['default'], ['default'], ['default'], ['default'], ['default']
+ return ['default'], ['csr'], ['csr'], ['csr'], ['csr']
+
def create_operator(self, ctx, shapes, dtypes):
return Sqr()
@@ -4064,14 +4085,100 @@ def create_operator(self, ctx, shapes, dtypes):
data = mx.symbol.cast(data, dtype='float64')
op = mx.symbol.cast(op, dtype='float32')
- x = mx.nd.array(np.random.uniform(-1, 1, size=(4, 10)))
- aux = mx.nd.zeros_like(x)
check_numeric_gradient(op, [x], [aux])
- x.attach_grad()
- with mx.contrib.autograd.train_section():
- y = mx.nd.Custom(x, aux, op_type='sqr')
+ data = mx.symbol.Variable('data', stype='csr')
+ aux = mx.symbol.Variable('aux')
+ op2 = mx.symbol.Custom(data=data, aux=aux, name='sqr', op_type='sqr')
+ x = x.tostype('csr')
+ aux = mx.nd.zeros_like(x)
+ check_numeric_gradient(op2, [x], [aux], grad_stype_dict={"data": "csr"})
+
+ x2 = mx.nd.array(np.random.uniform(-1, 1, size=(4, 10)))
+ x2 = x2.tostype('csr')
+ aux2 = mx.nd.zeros_like(x2)
+ x2.attach_grad()
+ with mx.autograd.record():
+ output = mx.nd.Custom(x2, aux2, name='sqr', op_type='sqr')
+ output.backward()
+ expected_output = mx.nd.sparse.square(x2)
+ expected_grad = 2 * x2
+ rtol = 1e-4
+ atol = 1e-6
+ assert_almost_equal(output.asnumpy(), expected_output.asnumpy(), rtol=rtol, atol=atol)
+ assert_almost_equal(x2.grad.asnumpy(), expected_grad.asnumpy(), rtol=rtol, atol=atol)
+
+
+ # test for backward compatibility, i.e. the correctness of default implementation of
+ # infer storage in custom operator
+ class Mult(mx.operator.CustomOp):
+ def forward(self, is_train, req, in_data, out_data, aux):
+ self.assign(out_data[0], req[0], in_data[0]*in_data[1])
+
+ def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
+ self.assign(in_grad[0], req[0], in_data[1])
+ self.assign(in_grad[1], req[1], in_data[0])
+
+ @mx.operator.register("mult")
+ class MultProp(mx.operator.CustomOpProp):
+ def __init__(self):
+ super(MultProp, self).__init__(need_top_grad=True)
+
+ def list_arguments(self):
+ return ['lhs', 'rhs']
+
+ def list_outputs(self):
+ return ['output']
+
+ def infer_shape(self, in_shape):
+ return in_shape, [in_shape[0]], []
+
+ def create_operator(self, ctx, shapes, dtypes):
+ return Mult()
+
+ lhs = mx.nd.array(np.random.uniform(-1, 1, size=(4, 10)))
+ rhs = mx.nd.array(np.random.uniform(-1, 1, size=(4, 10)))
+ lhs.attach_grad()
+ rhs.attach_grad()
+ with mx.autograd.record():
+ y = mx.nd.Custom(lhs, rhs, name='mult', op_type='mult')
y.backward()
+ assert_almost_equal(rhs.asnumpy(), lhs.grad.asnumpy(), rtol=rtol, atol=atol)
+ assert_almost_equal(lhs.asnumpy(), rhs.grad.asnumpy(), rtol=rtol, atol=atol)
+
+ class MultNoGrad(mx.operator.CustomOp):
+ def forward(self, is_train, req, in_data, out_data, aux):
+ self.assign(out_data[0], req[0], in_data[0]*in_data[1])
+
+ def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
+ self.assign(in_grad[0], req[0], in_data[1])
+ self.assign(in_grad[1], req[1], in_data[0])
+
+ @mx.operator.register("mult_no_grad")
+ class MultNoGradProp(mx.operator.CustomOpProp):
+ def __init__(self):
+ super(MultNoGradProp, self).__init__(need_top_grad=False)
+
+ def list_arguments(self):
+ return ['lhs', 'rhs']
+
+ def list_outputs(self):
+ return ['output']
+
+ def infer_shape(self, in_shape):
+ return in_shape, [in_shape[0]], []
+
+ def create_operator(self, ctx, shapes, dtypes):
+ return MultNoGrad()
+
+ def infer_storage_type_backward(self, ograd_stype, in_stype, out_stype, igrad_stype, aux_stype):
+ return ograd_stype, in_stype, out_stype, igrad_stype, aux_stype
+
+ with mx.autograd.record():
+ y2 = mx.nd.Custom(lhs, rhs, name="mult_no_grad", op_type="mult_no_grad")
+ y2.backward()
+ assert_almost_equal(rhs.asnumpy(), lhs.grad.asnumpy(), rtol=rtol, atol=atol)
+ assert_almost_equal(lhs.asnumpy(), rhs.grad.asnumpy(), rtol=rtol, atol=atol)
@with_seed()
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services