mxnet - 自定义 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

Operator 无法控制它将以何种模式执行。问题是，只有图优化器知道使用运算符的上下文，并且可以决定运算符是否需要在 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 并且运算符(operator)试图将差异添加到输出，而不是覆盖其中的内容，这将导致不可预测的结果，因为输出可能填充有垃圾。另一方面，如果模式是 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 的 future 版本中使用模式 kAddTo。同样在未来，可能会创建新的 API 来向图形优化器(或建议)发送提示以使用特定模式。但我不知道有这样的发展。

现在的问题是:“在哪种特定情况下 MXNet 0.10/0.11 将使用 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");