apeforest commented on a change in pull request #16218: Improving performance
of argmax operator
URL: https://github.com/apache/incubator-mxnet/pull/16218#discussion_r328384483
##########
File path: tests/python/unittest/test_operator.py
##########
@@ -7604,6 +7612,161 @@ def _softmaxoutput_normalization(multi_output,
use_ignore, normalization):
multi_output, use_ignore, normalization)
+@with_seed()
+def test_argmax():
+ ci_test = True # Use False if you need to collect aggregated
information about argmax performance
+ def getNumWorkers(shape, axis, ctx=mx.gpu()):
+ if ctx == mx.cpu():
+ return 1
+
+ nSteps = shape[axis]
+ nThreads = np.prod(shape)/nSteps
+ if nThreads > nSteps:
+ return 1
+
+ a = nSteps / nThreads
+ b = math.log2(a)
+ numbWorkers = math.floor(math.pow(2, (b * 5 + 28)/11))
+ return numbWorkers if 2 * numbWorkers <= nSteps else 1
+
+ def calc_argmax(tensor, axis, checkResult=False, nWorkers=1):
+ if not ci_test:
+ if axis is not None:
+ if nWorkers > 0:
+ # To use a predefined number of workers you need to
recompile MxNet with following
+ # lines added to broadcast_reduce_op.h:
+ #
+ # struct ReduceAxisParam : public
dmlc::Parameter<ReduceAxisParam> {
+ # ...
+ # int nWorkers;
+ # ...
+ # DMLC_DECLARE_FIELD(nWorkers).set_default(1) \
+ # .describe("Number of workers assigned for each
vector processing on GPU");
+ max = mx.nd.argmax(tensor, axis=axis, nWorkers=nWorkers)
+ else:
+ max = mx.nd.argmax(tensor, axis=axis, nWorkers=-1)
+ else:
+ max = mx.nd.argmax(tensor, axis=axis)
+
+ if checkResult:
+ topk_data = mx.nd.topk(tensor, axis=axis, is_ascend=0, k=1)
+ assert_almost_equal(max.reshape(-1).asnumpy(),
topk_data.reshape(-1).asnumpy())
+ else:
+ max.wait_to_read()
+ return max
+
+ # Useful function for collection of information about performance of
argmax and other MxNet operators
+ def runTest(a, b, shape = None, ctx=mx.gpu(), testArgmax=1,
checkResult=False, lenTest=1000, nWorkers=1):
+ debug = True #False
+ iFirst = 5 if lenTest > 5 else 0
+ iLast = lenTest - iFirst - 1
+ if iLast <= 0:
+ iLast = lenTest - 1
+
+ buff = None
+ jMax = 2 if shape is None and 2 * nWorkers <= a else 1
+ for j in range(jMax):
+ if buff is None:
+ if shape is None:
+ shape = (a, b) if j == 0 else (b, a)
+ if j == 0:
+ buff = "shape = (%2d, %6d): " % (a, b)
+ else:
+ buff = "{}".format(shape)
+
+ if testArgmax:
+ tmp = mx.nd.random.uniform(-1, 1, shape=shape, ctx=ctx)
+ else:
+ tmp = np.random.RandomState().uniform(-1, 1, shape)
+
+ for axis in range(len(shape)):
+ if nWorkers > shape[axis]:
+ buff += "***************"
+ continue
+ for i in range(lenTest):
+ if i == iFirst:
+ begin = time.time()
+ if i == iLast:
+ end = time.time()
+
+ if testArgmax:
+ max = calc_argmax(tmp, axis, checkResult=checkResult
and (i==0), nWorkers=nWorkers)
+ else:
+ max = np.argmax(tmp, axis=axis)
+
+ if max is not None and debug and testArgmax == 1:
+ if nWorkers < 0:
+ buff += " %6.3f (%3d)" % (end - begin,
getNumWorkers(shape, axis, ctx=ctx))
+ else:
+ buff += " %6.3f (%3d)" % (end - begin, nWorkers)
+ else:
+ buff += " %6.3f" % (end - begin)
+
+ if j == jMax - 1:
+ if nWorkers > 0:
+ buff += " ** %3d" % getNumWorkers(shape, axis, ctx=ctx)
+ print(buff)
Review comment:
Please do not add print in unit test
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
