Custom MXNet and kAddTo Operators
I am writing a custom C ++ operator in MXNet and I am having a hard time finding documentation when kAddTo
set in an operator call. As a minimal example, let's say my new operator is called foo()
and I want to do the following calculation:
A = mx.sym.Variable('A')
B = mx.sym.Variable('B')
T = mx.sym.foo(A)
T += mx.sym.foo(B)
In general, how can I ensure that the fourth row above accumulates in T, rather than creating a new temporary storage for the result mx.sym.foo(B)
, and then doing the computation T = T + temp
?
(Using the Kernighan-Ritchie debugger as with the print instructions, I found myself kWriteTo
set on both lines three and four. The enumeration is kAddTo
never set.)
A little more detail on my specific problem: in my current implementation, it foo()
zeroes out the output memory before performing a computation that fills it with the appropriate values. I definitely only want to perform this zeroing when creating a new output location, not when accumulating in an existing one.
Update
Offline, a colleague suggested using
mx.sym.elemwise_add(lhs=T, rhs=mx.sym.foo(B), out=T)
instead of line 4 above. However, I could still see what was getting kWriteTo
set on both lines of the calculation. Then I got the following response:
"Memory scheduling and in-place operations are done automatically. This will be done automatically. Users need not worry about this." which probably means it is
req[0]
not an accurate indicator in this case.If you want to check if its addplace inplace is there, you can print the valueoutputs[0].dptr_
andlhs.dptr_
to see if they are equal.
I have not tested this yet.
source to share
The operator cannot control in which mode it will be executed. The point is that only the graph optimizer knows the context in which the operator is used and can decide if the operator should be executed in kWriteTo
or kAddTo
. More precisely, it happens here in the DetectInplaceAddTo method . And even if in some cases this was done in kAddTo
, this behavior may be changed in the future due to changes in the logic that optimizes the computational graph.
"Memory scheduling and in-place operations are done automatically. This will be done automatically. Users do not need to worry about that."
This means that the operator has no control over which mode it is executed in, however the operator MUST strictly adhere to the mode that was requested ( kWriteTo
or kAddTo
). For example, if a mode kWriteTo
, and a statement tries to add diff to the outputs, instead of overriding what is in it, it will lead to unpredictable results, as the outputs can be filled with garbage. On the other hand, if the mode, kAddTo
however, the operator does not support it, it can be even worse, since instead of adding the results to the outputs, it simply overrides the outputs (such cases are usually very difficult to debug), This sometimes leads to errors like this .
So, in short:
In general, how can I ensure that the fourth row above accumulates into T, as opposed to creating a new temporary storage for the result of mx.sym.foo (B) and then doing the computation T = T + temp?
You cannot, it is not a decision of the operator in which the mode should be executed. Even if the configuration uses kAddTo
future MXNet version mode. It is also possible in the future to create new APIs to send a hint to the chart optimizer (or proposal) for using a certain mode. But I am not aware of this development.
Now the question is: "In what specific case will MXNet 0.10 / 0.11 use kAddTo
"?
It's tricky looking at the following code :
for (uint32_t nid = 0; nid < idx.num_nodes(); ++nid) {
const auto& inode = idx[nid];
if (inode.source->op() != ewise_plus_op) continue; // <= HERE
int sid = storage_id[idx.entry_id(inode.inputs[0])];
It looks like it's kAddTo
only used during time _grad_add
, which is sad. It could also be a bug, as perhaps instead of:
static const Op* ewise_plus_op = Op::Get("_grad_add");
The actual intention was:
static const Op* ewise_plus_op = Op::Get("elemwise_add");
source to share