From 3cd1d0197aed867e72ab0974b253fd8d815ee33e Mon Sep 17 00:00:00 2001 From: Ciyong Chen Date: Mon, 13 Aug 2018 16:33:34 +0800 Subject: [PATCH] Change to use MXNet's topk for CPUs in inference Since MXNet's topk has better performance than numpy version with PR https://github.com/apache/incubator-mxnet/pull/12085, in order to leverage such performance boost, change to use MXNet's topk for CPU device when doing inference. --- sockeye/inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sockeye/inference.py b/sockeye/inference.py index bdec9f468..e42f7e7fb 100644 --- a/sockeye/inference.py +++ b/sockeye/inference.py @@ -1021,7 +1021,7 @@ def __init__(self, k=self.beam_size, batch_size=self.batch_size, offset=self.offset, - use_mxnet_topk=self.context != mx.cpu()) # MXNet implementation is faster on GPUs + use_mxnet_topk=True) self._sort_by_index = SortByIndex() self._sort_by_index.initialize(ctx=self.context)