Add int64 support to sparse_to_dense_mask_op
Summary: [CAFFE2] Add int64 support to sparse_to_dense_mask_op
Reviewed By: ender-wieczorek
Differential Revision: D6022278
fbshipit-source-id: 489b6df4d43a64c743ee278d94929ca50259f7b8
diff --git a/caffe2/operators/sparse_to_dense_mask_op.cc b/caffe2/operators/sparse_to_dense_mask_op.cc
index d84553f..f234885 100644
--- a/caffe2/operators/sparse_to_dense_mask_op.cc
+++ b/caffe2/operators/sparse_to_dense_mask_op.cc
@@ -30,7 +30,7 @@
.TensorInferenceFunction([](const OperatorDef& def,
const vector<TensorShape>& in) {
ArgumentHelper helper(def);
- auto mask = helper.template GetRepeatedArgument<int>("mask");
+ auto mask = helper.template GetRepeatedArgument<int64_t>("mask");
bool return_presence_mask = helper.template GetSingleArgument<bool>(
"return_presence_mask", false);
vector<TensorShape> out(1);
diff --git a/caffe2/operators/sparse_to_dense_mask_op.h b/caffe2/operators/sparse_to_dense_mask_op.h
index f508a01..94af59d 100644
--- a/caffe2/operators/sparse_to_dense_mask_op.h
+++ b/caffe2/operators/sparse_to_dense_mask_op.h
@@ -33,14 +33,15 @@
USE_OPERATOR_CONTEXT_FUNCTIONS;
SparseToDenseMaskBase(const OperatorDef& operator_def, Workspace* ws)
: Operator<Context>(operator_def, ws) {
- std::vector<int> mask =
- OperatorBase::template GetRepeatedArgument<int>("mask");
+ std::vector<int64_t> mask =
+ OperatorBase::template GetRepeatedArgument<int64_t>("mask");
featuresCount_ = mask.size();
+
CAFFE_ENFORCE(!mask.empty(), "mask can't be empty");
auto biggest = *std::max_element(mask.begin(), mask.end());
dense_.assign(std::min(kMaxDenseSize, biggest + 1), -1);
for (int i = 0; i < mask.size(); i++) {
- int id = mask[i];
+ int64_t id = mask[i];
CAFFE_ENFORCE_GE(id, 0, "Only positive IDs are allowed.");
if (id >= kMaxDenseSize) {
CAFFE_ENFORCE(sparse_.count(id) == 0, "Duplicated id: ", id);
@@ -53,13 +54,13 @@
}
protected:
- const int kMaxDenseSize = 1024 * 128;
+ const int64_t kMaxDenseSize = 1024 * 128;
- std::unordered_map<int, int> sparse_;
+ std::unordered_map<int64_t, int> sparse_;
std::vector<int> dense_;
int featuresCount_;
- inline int getFeatureIdx(int id) const {
+ inline int getFeatureIdx(int64_t id) const {
if (id >= kMaxDenseSize) {
const auto& iter = sparse_.find(id);
if (iter == sparse_.end()) {
@@ -160,7 +161,7 @@
rows * cols, false, presence_mask_data, &context_);
}
- int32_t offset = 0;
+ int64_t offset = 0;
for (int r = 0; r < rows; r++) {
for (int c = 0; c < lengths_vec[r]; c++) {
const auto sparse_index = sparse_indices_vec[offset + c];
diff --git a/caffe2/python/operator_test/sparse_to_dense_mask_op_test.py b/caffe2/python/operator_test/sparse_to_dense_mask_op_test.py
index 3d774c4..cdcdfad 100644
--- a/caffe2/python/operator_test/sparse_to_dense_mask_op_test.py
+++ b/caffe2/python/operator_test/sparse_to_dense_mask_op_test.py
@@ -60,6 +60,39 @@
gc, op, input_data, 1, [0])
@given(n=st.integers(1, 10), k=st.integers(1, 5),
+ use_length=st.booleans(), **hu.gcs_cpu_only)
+ def test_sparse_to_dense_mask_with_int64(self, n, k, use_length, gc, dc):
+ lengths = np.random.randint(k, size=n).astype(np.int32) + 1
+ N = sum(lengths)
+ int64_mask = 10000000000
+ indices = np.random.randint(5, size=N) + int64_mask
+ values = np.random.rand(N, 2).astype(np.float32)
+ default = np.random.rand(2).astype(np.float32)
+ mask = np.arange(3) + int64_mask
+ np.random.shuffle(mask)
+
+ input_str = ['indices', 'values', 'default']
+ input_data = [indices, values, default]
+ if use_length and n > 1:
+ input_str.append('lengths')
+ input_data.append(lengths)
+ output_str = ['output']
+
+ op = core.CreateOperator(
+ 'SparseToDenseMask',
+ input_str,
+ output_str,
+ mask=mask,
+ )
+
+ # Check over multiple devices
+ self.assertDeviceChecks(
+ dc, op, input_data, [0])
+ # Gradient check for values
+ self.assertGradientChecks(
+ gc, op, input_data, 1, [0])
+
+ @given(n=st.integers(1, 10), k=st.integers(1, 5),
dim=st.integers(1, 3), **hu.gcs_cpu_only)
def test_sparse_to_dense_mask_high_dim(self, n, k, dim, gc, dc):
lengths = np.random.randint(k, size=n).astype(np.int32) + 1
