自定义 MXNet 运算符和 kAddTo

Question

我正在 MXNet 中编写 C++ 自定义运算符，但无法找到有关何时在运算符调用中设置 kAddTo 的文档。作为一个最小的例子，假设我的新运算符称为 foo()，我想执行以下计算：

A = mx.sym.Variable('A')
B = mx.sym.Variable('B')
T = mx.sym.foo(A)
T += mx.sym.foo(B)

一般来说，如何确保上面的第四行累加到T而不是为mx.sym.foo(B)的结果创建一个新的临时存储然后执行T = T + temp计算？

（使用 Kernighan-Ritchie 调试器，又名打印语句，我发现 kWriteTo 在第三行和第四行都设置了。枚举 kAddTo 从未设置。 )

关于我的具体问题的更多细节：在我当前的实现中，foo()在执行计算之前将输出内存清零，并用适当的值填充它。我绝对只想在创建新输出位置时执行此清零，而不是在累积到现有位置时。

更新

离线，同事建议使用

mx.sym.elemwise_add(lhs=T, rhs=mx.sym.foo(B), out=T)

代替上面的第 4 行。但是，我仍然看到 kWriteTo 在两行计算中都被设置了。然后我收到了以下回复：

“Memory planning and inplace operations are automatic. It will be done automatically. Users don’t need to worry about it.”, which probably means that req[0] is not an accurate indicator in this case. If you want to verify whether it’s an inplace addTo, you can print out the value of outputs[0].dptr_ and lhs.dptr_ to see whether they are equal.

我还没有检查过这个。

Answer 1

操作员无法控制它将在哪种模式下执行。问题是，只有图优化器知道使用运算符的上下文，并且可以决定是否需要在 kWriteTo 或 kAddTo 中执行运算符。更准确地说，这种情况会发生 here in the method DetectInplaceAddTo。即使在某些情况下它已在 kAddTo 中执行，但由于优化计算图的逻辑发生变化，这种行为将来可能会发生变化。

“Memory planning and inplace operations are automatic. It will be done automatically. Users don’t need to worry about it.”

这意味着操作员无法控制它在哪种模式下执行，但是操作员必须严格遵守已请求的模式（kWriteTo 或 kAddTo）。例如，如果模式是 kWriteTo 并且操作员试图将差异添加到输出，而不是覆盖其中的内容，这将导致不可预测的结果，因为输出可能填充了垃圾。另一方面，如果模式是 kAddTo 但是操作员不支持它可能会更糟，因为它不会将结果添加到输出，而是会覆盖输出（像这样的情况通常很难调试）。这有时会导致 this one.

之类的错误

所以，简而言之：

In general, how do I ensure that the fourth line above accumulates into T as opposed to creating a new temporary storage for the result of mx.sym.foo(B) and then performing the T = T + temp calculation?

你不能，这不是运营商决定在哪种模式下执行。即使配置在 MXNet 的未来版本中使用模式 kAddTo。同样在未来，可能会创建新的 API 来向图形优化器（或建议）发送提示以使用特定模式。但是我不知道有这样的发展。

现在问题："in which particular case MXNet 0.10/0.11 will use kAddTo"?

这很棘手，通过查看 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])];

看来kAddTo只在_grad_add期间用过，可惜了。这也可能是一个错误，因为也许不是：

static const Op* ewise_plus_op = Op::Get("_grad_add");

实际意图是：

static const Op* ewise_plus_op = Op::Get("elemwise_add");

自定义 MXNet 运算符和 kAddTo

Custom MXNet Operators and kAddTo

mxnet