1 #ifndef CAFFE2_OPERATORS_SPARSE_TO_DENSE_MASK_OP_H_ 2 #define CAFFE2_OPERATORS_SPARSE_TO_DENSE_MASK_OP_H_ 5 #include <unordered_map> 7 #include "caffe2/core/context.h" 8 #include "caffe2/core/operator.h" 9 #include "caffe2/core/tensor.h" 10 #include "caffe2/utils/math.h" 14 template <
class Context>
17 USE_OPERATOR_CONTEXT_FUNCTIONS;
20 std::vector<int64_t> mask =
21 OperatorBase::template GetRepeatedArgument<int64_t>(
"mask");
22 featuresCount_ = mask.size();
24 CAFFE_ENFORCE(!mask.empty(),
"mask can't be empty");
25 auto biggest = *std::max_element(mask.begin(), mask.end());
26 dense_.assign(std::min(kMaxDenseSize, biggest + 1), -1);
27 for (
int i = 0; i < mask.size(); i++) {
29 CAFFE_ENFORCE_GE(
id, 0,
"Only positive IDs are allowed.");
30 if (
id >= kMaxDenseSize) {
31 CAFFE_ENFORCE(sparse_.count(
id) == 0,
"Duplicated id: ", id);
34 CAFFE_ENFORCE(dense_[
id] == -1,
"Duplicated id: ",
id);
41 const int64_t kMaxDenseSize = 1024 * 128;
43 std::unordered_map<int64_t, int> sparse_;
44 std::vector<int> dense_;
47 inline int getFeatureIdx(int64_t
id)
const {
48 if (
id >= kMaxDenseSize) {
49 const auto& iter = sparse_.find(
id);
50 if (iter == sparse_.end()) {
56 return (
id >= dense_.size()) ? -1 : dense_[
id];
61 template <
class Context>
64 USE_OPERATOR_CONTEXT_FUNCTIONS;
67 returnPresenceMask_ = OperatorBase::template GetSingleArgument<bool>(
68 "return_presence_mask",
false);
69 maxSkippedSparseIndices_ =
70 OperatorBase::template GetSingleArgument<int32_t>(
71 "max_skipped_indices", kMaxSkippedSparseIndices);
74 bool RunOnDevice()
override {
76 this, Input(INDICES));
79 template <
typename TInd>
80 bool DoRunWithType() {
81 auto& sparse_indices = Input(INDICES);
82 CAFFE_ENFORCE_EQ(sparse_indices.ndim(), 1);
83 auto& sparse_values = Input(VALUES);
84 CAFFE_ENFORCE_GE(sparse_values.ndim(), 1);
85 CAFFE_ENFORCE_EQ(sparse_indices.size(), sparse_values.dim(0));
86 auto& default_value = Input(DEFAULT);
87 CAFFE_ENFORCE_EQ(default_value.ndim() + 1, sparse_values.ndim());
88 CAFFE_ENFORCE_EQ(default_value.size(), sparse_values.size_from_dim(1));
89 CAFFE_ENFORCE(sparse_values.meta() == default_value.meta());
91 const TInd* sparse_indices_vec = sparse_indices.template data<TInd>();
92 const char* sparse_values_vec =
93 static_cast<const char*
>(sparse_values.raw_data());
94 const void* default_val = default_value.raw_data();
96 TIndex block_size = default_value.size();
97 size_t block_nbytes = default_value.nbytes();
99 const int cols = this->featuresCount_;
101 int32_t sparse_indices_length = sparse_indices.dim32(0);
102 const int32_t* lengths_vec =
nullptr;
103 auto* output = Output(OUTPUTVALUE);
105 if (returnPresenceMask_) {
106 presence_mask = Output(PRESENCEMASK);
108 vector<TIndex> shape;
109 if (InputSize() == 4) {
110 auto& lengths = Input(LENGTHS);
111 CAFFE_ENFORCE_EQ(lengths.ndim(), 1);
112 lengths_vec = lengths.template data<int32_t>();
113 rows = lengths.dim32(0);
118 lengths_vec = &sparse_indices_length;
120 shape.push_back(rows);
122 shape.push_back(cols);
123 if (returnPresenceMask_) {
124 presence_mask->
Resize(shape);
127 shape.end(), default_value.dims().begin(), default_value.dims().end());
128 output->Resize(shape);
133 static_cast<char*
>(output->raw_mutable_data(sparse_values.meta()));
134 for (
int i = 0; i < cols * rows; i++) {
135 context_.template CopyItems<Context, Context>(
136 default_value.meta(),
139 output_data + i * block_nbytes);
141 bool* presence_mask_data =
nullptr;
142 if (returnPresenceMask_) {
143 presence_mask_data = presence_mask->template mutable_data<bool>();
144 math::Set<bool, Context>(
145 rows * cols,
false, presence_mask_data, &context_);
149 for (
int r = 0; r < rows; r++) {
150 for (
int c = 0; c < lengths_vec[r]; c++) {
151 const auto sparse_index = sparse_indices_vec[offset + c];
152 if (sparse_index < 0 ||
153 sparse_index >= std::numeric_limits<TInd>::max()) {
155 ++skippedSparseIndices_,
156 maxSkippedSparseIndices_,
157 "Too many sparse indices skipped");
160 int idx = this->getFeatureIdx(sparse_index);
162 context_.template CopyItems<Context, Context>(
163 sparse_values.meta(),
165 sparse_values_vec + (offset + c) * block_nbytes,
166 output_data + (r * cols + idx) * block_nbytes);
167 if (returnPresenceMask_) {
168 presence_mask_data[r * cols + idx] =
true;
172 offset += lengths_vec[r];
179 static const uint32_t kMaxSkippedSparseIndices = 5;
181 bool returnPresenceMask_;
182 uint32_t maxSkippedSparseIndices_ = 0;
183 uint32_t skippedSparseIndices_ = 0;
185 INPUT_TAGS(INDICES, VALUES, DEFAULT, LENGTHS);
186 OUTPUT_TAGS(OUTPUTVALUE, PRESENCEMASK);
189 template <
class Context>
192 USE_OPERATOR_CONTEXT_FUNCTIONS;
196 bool RunOnDevice()
override {
198 this, Input(INDICES));
201 template <
typename TInd>
202 bool DoRunWithType() {
203 auto& sparse_indices = Input(INDICES);
204 CAFFE_ENFORCE_EQ(sparse_indices.ndim(), 1);
205 auto& gradient_output = Input(GOUTPUT);
207 TIndex block_size = gradient_output.size_from_dim(1);
208 size_t block_nbytes = gradient_output.itemsize() * block_size;
210 const int cols = this->featuresCount_;
213 int32_t default_length = sparse_indices.dim32(0);
214 const int32_t* lengths_vec =
nullptr;
215 auto* output = Output(GVALUES);
216 vector<TIndex> shape;
217 if (InputSize() > LENGTHS) {
220 auto& lengths = Input(LENGTHS);
221 lengths_vec = lengths.template data<int32_t>();
222 rows = lengths.dim32(0);
223 CAFFE_ENFORCE_EQ(lengths.ndim(), 1);
224 CAFFE_ENFORCE_GE(gradient_output.ndim(), 2);
225 CAFFE_ENFORCE_EQ(gradient_output.dim(0), rows);
226 CAFFE_ENFORCE_EQ(gradient_output.dim(1), cols);
227 block_nbytes /= gradient_output.dim(1);
228 block_size /= gradient_output.dim(1);
235 lengths_vec = &default_length;
236 CAFFE_ENFORCE_GE(gradient_output.ndim(), 1);
237 CAFFE_ENFORCE_EQ(gradient_output.dim(0), cols);
239 shape.push_back(default_length);
243 gradient_output.dims().begin() + iter_offset,
244 gradient_output.dims().end());
245 output->Resize(shape);
247 const TInd* sparse_indices_vec = sparse_indices.template data<TInd>();
248 const char* gradient_output_vec =
249 static_cast<const char*
>(gradient_output.raw_data());
252 static_cast<char*
>(output->raw_mutable_data(gradient_output.meta()));
253 math::Set<char, Context>(
254 default_length * gradient_output.itemsize(), 0, output_data, &context_);
259 vector<bool> gradient_used(cols,
false);
260 for (
int r = 0; r < rows; r++) {
261 std::fill(gradient_used.begin(), gradient_used.end(),
false);
262 for (
int c = lengths_vec[r] - 1; c >= 0; c--) {
263 int idx = this->getFeatureIdx(sparse_indices_vec[offset + c]);
264 if (idx != -1 && !gradient_used[idx]) {
265 gradient_used[idx] =
true;
266 context_.template CopyItems<Context, Context>(
267 gradient_output.meta(),
269 gradient_output_vec + (r * cols + idx) * block_nbytes,
270 output_data + (offset + c) * block_nbytes);
273 offset += lengths_vec[r];
279 INPUT_TAGS(INDICES, GOUTPUT, LENGTHS);
280 OUTPUT_TAGS(GVALUES);
285 #endif // CAFFE2_OPERATORS_SPARSE_TO_DENSE_MASK_OP_H_
Tensor is the basic class in Caffe2 that stores a contiguous memory with its shape information...
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
void Resize(Ts...dim_source)
Resizes a tensor.
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...