使用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