使用 topk 修剪 MXNet 中的输入

Question

我正在尝试在 MXNet 中创建一个运算符，它将通过以下方式在输出中引入稀疏性：

分别对每个数据点进行修剪（轴 0 用于数据点）
将较低的权重降至 0
保持与输入相同的维度

我目前正在使用以下代码执行此操作（假设 act 是此运算符的输入）：

flat = mx.sym.flatten(act)
mask = mx.sym.topk(flat, k = int(frac * flat.infer_shape(data=shape)[1][0][1]), axis = 1, ret_typ = 'mask').reshape(act.infer_shape(data=shape)[1][0])
custom = mx.sym.where(mask == 1, act, mask)

在这个实现中，张量行为的维数有限制。一个非常大的张量，当展平并传递到 topk 时会导致 IndexFill 错误：

[20:27:53] /home/ubuntu/mxnet/dmlc-core/include/dmlc/logging.h:304: [20:27:53] /home/ubuntu/mxnet/mshadow/mshadow/././././cuda/tensor_gpu-inl.cuh:58: too large launch parameter: IndexFill[100352,1], [32,32,1]

Stack trace returned 10 entries:
[bt] (0) /usr/local/lib/python2.7/dist-packages/mxnet-0.10.1-py2.7.egg/mxnet/libmxnet.so(_ZN4dmlc15LogMessageFatalD1Ev+0x3c) [0x7fb593bbc9ac]
[bt] (1) /usr/local/lib/python2.7/dist-packages/mxnet-0.10.1-py2.7.egg/mxnet/libmxnet.so(_ZN7mshadow4cuda9IndexFillIffEEvNS_6TensorINS_3gpuELi2ET0_EERKNS2_IS3_Li1ET_EERKS5_+0x492) [0x7fb59581bf82]
[bt] (2) /usr/local/lib/python2.7/dist-packages/mxnet-0.10.1-py2.7.egg/mxnet/libmxnet.so(_ZN5mxnet2op8TopKImplIN7mshadow3gpuEEEvNS_10RunContextENS_8ResourceERKNS_5TBlobERKSt6vectorIS6_SaIS6_EERKNS0_9TopKParamE+0x3ca1) [0x7fb595841521]
[bt] (3) /usr/local/lib/python2.7/dist-packages/mxnet-0.10.1-py2.7.egg/mxnet/libmxnet.so(_ZN5mxnet2op4TopKIN7mshadow3gpuEEEvRKN4nnvm9NodeAttrsERKNS_9OpContextERKSt6vectorINS_5TBlobESaISC_EERKSB_INS_9OpReqTypeESaISH_EESG_+0x345) [0x7fb595842cc5]
[bt] (4) /usr/local/lib/python2.7/dist-packages/mxnet-0.10.1-py2.7.egg/mxnet/libmxnet.so(+0x1318cf9) [0x7fb5947aecf9]
[bt] (5) /usr/local/lib/python2.7/dist-packages/mxnet-0.10.1-py2.7.egg/mxnet/libmxnet.so(_ZN5mxnet6engine14ThreadedEngine15ExecuteOprBlockENS_10RunContextEPNS0_8OprBlockE+0x8c) [0x7fb5947ef07c]
[bt] (6) /usr/local/lib/python2.7/dist-packages/mxnet-0.10.1-py2.7.egg/mxnet/libmxnet.so(_ZNSt17_Function_handlerIFvvEZZN5mxnet6engine23ThreadedEnginePerDevice13PushToExecuteEPNS2_8OprBlockEbENKUlvE1_clEvEUlvE_E9_M_invokeERKSt9_Any_data+0x60) [0x7fb5947f2190]
[bt] (7) /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xb1a60) [0x7fb5a3c45a60]
[bt] (8) /lib/x86_64-linux-gnu/libpthread.so.0(+0x8184) [0x7fb5a9e07184]
[bt] (9) /lib/x86_64-linux-gnu/libc.so.6(clone+0x6d) [0x7fb5a9b34bed]

所以我的问题是：

它目前以非常小的批量运行。但是有没有办法增加批量大小并避免错误？
有没有更好的实现操作符的方法？

Answer 1

问题的原因与 GPU 运算符及其内核的实现有关，特别是线程数、块数，以及内核启动时的网格尺寸。

特别是，NVIDIA CUDA 计算能力指定了最大线程数、每个块的线程数以及每个维度（网格维度）的块数。例如，参见 http://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capabilities.

在您的情况下，第一个网格维度超过了 65535 的阈值。在 MXNet 中，这个阈值也被定义为 kMaxGridDim。因此它抛出错误。

要解决这个问题，可能有不同的选择：更改特定的运算符本身和内核启动请求的线程数，然后可能更改内核本身；或者，修复通用 MXNet GPU 内核启动函数也可以解决问题。

明天我会调查一下，问题解决后更新我的答案。

编辑：问题已得到解决：https://github.com/dmlc/mshadow/pull/277

使用 topk 修剪 MXNet 中的输入

Using topk to prune inputs in MXNet

python

mxnet