This is an automated email from the ASF dual-hosted git repository.
apeforest pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new d51753b [Large Tensor] Fixed Embedding op (#17599)
d51753b is described below
commit d51753bd87fb24509c01b7e9f8f638dd1026c515
Author: Connor Goggins <cgoggi...@gmail.com>
AuthorDate: Mon Feb 24 23:33:10 2020 -0800
[Large Tensor] Fixed Embedding op (#17599)
* Switched from int to index_t for input_dim
* Implemented fix for output_dim
* Added nightly test for Embedding
* Set const value for output dim
* More standardization via const param
---
src/operator/tensor/indexing_op.h | 8 ++++----
tests/nightly/test_large_array.py | 13 +++++++++++++
2 files changed, 17 insertions(+), 4 deletions(-)
diff --git a/src/operator/tensor/indexing_op.h
b/src/operator/tensor/indexing_op.h
index 7b6c16a..5449fbe 100644
--- a/src/operator/tensor/indexing_op.h
+++ b/src/operator/tensor/indexing_op.h
@@ -66,8 +66,8 @@ enum QuantizedEmbeddingOpResource {kTempSpace};
struct SparseEmbeddingParam: public dmlc::Parameter<SparseEmbeddingParam> {
- int input_dim;
- int output_dim;
+ index_t input_dim;
+ index_t output_dim;
int dtype;
bool deterministic;
DMLC_DECLARE_PARAMETER(SparseEmbeddingParam) {
@@ -89,8 +89,8 @@ struct SparseEmbeddingParam: public
dmlc::Parameter<SparseEmbeddingParam> {
};
struct EmbeddingParam: public dmlc::Parameter<EmbeddingParam> {
- int input_dim;
- int output_dim;
+ index_t input_dim;
+ index_t output_dim;
int dtype;
bool sparse_grad;
DMLC_DECLARE_PARAMETER(EmbeddingParam) {
diff --git a/tests/nightly/test_large_array.py
b/tests/nightly/test_large_array.py
index edf796c..8b36d09 100644
--- a/tests/nightly/test_large_array.py
+++ b/tests/nightly/test_large_array.py
@@ -38,6 +38,7 @@ LARGE_X = 100000000
SMALL_X = 100
SMALL_Y = 50
LARGE_SIZE = LARGE_X * SMALL_Y
+LARGE_TENSOR_SHAPE = 2**32
def test_nn():
@@ -467,6 +468,17 @@ def test_nn():
assert res.shape[2] == 2
assert res.shape[3] == 2
assert res.shape[4] == 1
+ def check_embedding():
+ data = nd.random_normal(shape=(LARGE_TENSOR_SHAPE, 1))
+ weight = nd.random_normal(shape=(LARGE_TENSOR_SHAPE, 1))
+ input_dim = LARGE_TENSOR_SHAPE
+ output_dim = 1
+
+ out = nd.Embedding(data=data, weight=weight, input_dim=input_dim,
output_dim=output_dim)
+
+ assert out.shape[0] == LARGE_TENSOR_SHAPE
+ assert out.shape[1] == 1
+ assert out.shape[2] == 1
check_gluon_embedding()
check_fully_connected()
@@ -488,6 +500,7 @@ def test_nn():
check_l2_normalization()
check_instance_norm()
check_col2im()
+ check_embedding()
def test_tensor():
