1 #include "caffe2/operators/one_hot_ops.h" 3 #include "caffe2/core/operator.h" 4 #include "caffe2/core/tensor.h" 10 bool BatchOneHotOp<CPUContext>::DoRunWithType() {
11 auto& input = Input(X);
12 auto& lens = Input(LENS);
13 auto& vals = Input(VALS);
14 CAFFE_ENFORCE_GE(input.ndim(), 1);
15 auto N = input.dim(0);
16 auto D = input.size_from_dim(1);
17 CAFFE_ENFORCE_EQ(lens.size(), D);
19 const auto* lens_data = lens.template data<int32_t>();
20 TIndex output_dim = 0;
21 valsOffsets_.resize(D + 1);
22 for (TIndex i = 0; i < D; i++) {
23 CAFFE_ENFORCE_GE(lens_data[i], 0);
24 valsOffsets_[i] = output_dim;
25 output_dim += lens_data[i];
27 valsOffsets_[D] = output_dim;
29 CAFFE_ENFORCE_EQ(vals.size(), output_dim);
30 auto* output = Output(ONE_HOT);
31 output->Resize(N, output_dim);
33 const auto* input_data = input.template data<T>();
34 const auto* vals_data = vals.template data<T>();
35 auto* output_data = output->template mutable_data<T>();
37 for (TIndex i = 0; i < N; ++i) {
38 for (TIndex j = 0; j < D; j++) {
39 const auto input_val = input_data[i * D + j];
40 for (TIndex k = valsOffsets_[j]; k < valsOffsets_[j + 1]; ++k) {
41 output_data[k] = vals_data[k] == input_val;
44 output_data += output_dim;
50 vector<TensorShape> TensorInferenceForBatchOneHot(
52 const vector<TensorShape>& in) {
53 std::vector<TIndex> output_dims(2);
54 output_dims[0] = in[0].dims(0);
55 output_dims[1] = in[2].dims(0);
56 return vector<TensorShape>{
57 CreateTensorShape(vector<TIndex>{output_dims}, in[0].data_type())};
60 OpSchema::Cost CostInferenceForBatchOneHot(
61 const OperatorDef& def,
62 const vector<TensorShape>& in) {
63 struct OpSchema::Cost c;
64 const TensorShape output = TensorInferenceForBatchOneHot(def, in)[0];
67 c.bytes_moved = output.dims(0) * output.dims(1) *
sizeof(int32_t);
73 void OneHotOp<CPUContext>::DoOneHotOp(
78 const TIndex* indices_ptr = indices.template data<TIndex>();
79 float* one_hots_ptr = one_hots->template mutable_data<float>();
80 memset(one_hots_ptr, 0, one_hots->nbytes());
81 for (
int i = 0; i < batch_size; ++i) {
82 auto label_idx = indices_ptr[i];
83 DCHECK((0 <= label_idx) && (label_idx < index_size));
84 one_hots_ptr[label_idx] = 1.0;
85 one_hots_ptr += index_size;
90 bool BatchBucketOneHotOp<CPUContext>::RunOnDevice() {
91 auto& input = Input(X);
92 auto& lens = Input(LENS);
93 auto& boundaries = Input(BOUNDARIES);
94 CAFFE_ENFORCE_GE(input.ndim(), 1);
95 auto N = input.dim(0);
96 auto D = input.size_from_dim(1);
97 CAFFE_ENFORCE_EQ(lens.size(), D);
99 const auto* lens_data = lens.template data<int32_t>();
102 std::accumulate(lens_data, lens_data + lens.size(), 0),
104 "The sum of length should be equal to the length of boundaries");
106 TIndex output_dim = 0;
107 for (TIndex i = 0; i < D; i++) {
108 CAFFE_ENFORCE_GT(lens_data[i], 0);
110 output_dim += (lens_data[i] + 1);
112 auto* output = Output(ONE_HOT);
113 output->Resize(N, output_dim);
115 const auto* input_data = input.template data<float>();
116 const auto* boundaries_data = boundaries.template data<float>();
117 auto* output_data = output->template mutable_data<float>();
119 math::Set<float, CPUContext>(output->size(), 0.f, output_data, &context_);
122 for (TIndex i = 0; i < N; i++) {
123 auto* boundaries_offset = boundaries_data;
124 TIndex output_offset = 0;
126 for (TIndex j = 0; j < D; j++) {
128 TIndex bucket_idx = std::lower_bound(
130 boundaries_offset + lens_data[j],
133 output_data[i * output_dim + output_offset + bucket_idx] = 1.0;
134 boundaries_offset += lens_data[j];
135 output_offset += (lens_data[j] + 1);
148 bool RunOnDevice()
override {
149 auto& lengths = Input(0);
150 auto& indices = Input(1);
151 auto& index_size_tensor = Input(2);
152 CAFFE_ENFORCE(lengths.ndim() == 1);
153 CAFFE_ENFORCE(indices.ndim() == 1);
154 CAFFE_ENFORCE(index_size_tensor.size() == 1);
155 auto batch_size = lengths.size();
156 auto index_size = *index_size_tensor.data<int64_t>();
157 CAFFE_ENFORCE(index_size > 0);
159 auto* lengths_ptr = lengths.data<int32_t>();
160 auto* indices_ptr = indices.data<int64_t>();
161 auto* one_hots = Output(0);
162 one_hots->Resize(batch_size, index_size);
163 auto* one_hots_ptr = one_hots->mutable_data<
float>();
164 if (one_hots->size() == 0) {
167 memset(one_hots_ptr, 0, one_hots->nbytes());
169 for (
int i = 0; i < batch_size; ++i) {
170 for (
int j = 0; j < lengths_ptr[i]; ++j) {
171 DCHECK(el_idx < indices.size());
172 auto label_idx = indices_ptr[el_idx++];
173 DCHECK((0 <= label_idx) && (label_idx < index_size));
174 one_hots_ptr[label_idx] = 1.0;
176 one_hots_ptr += index_size;
186 OPERATOR_SCHEMA(BatchBucketOneHot)
190 Input is a matrix tensor. Its first dimension is the batch 191 size. For each column, bucketize it based on the boundary values and then do 192 one hot encoding. The `lengths` specifies the number of boundary values for each 193 column. The final number of buckets is this number plus 1. This would also be 194 the expanded feature size. `boundaries` specifies all the boundary values. 195 Note that each bucket is right-inclusive. That is, given boundary values 196 [b1, b2, b3], the buckets are defined as (-int, b1], (b1, b2], (b2, b3], (b3, inf). 199 If data = [[2, 3], [4, 1], [2, 5]], lengths = [2, 3], 200 and boundaries = [0.1, 2.5, 1, 3.1, 4.5], then 202 output = [[0, 1, 0, 0, 1, 0, 0], [0, 0, 1, 1, 0, 0, 0], [0, 1, 0, 0, 0, 0, 1]] 205 .Input(0, "data",
"input tensor matrix")
206 .Input(1,
"lengths",
"the size is the same as the width of the `data`")
207 .Input(2,
"boundaries",
"bucket boundaries")
211 "output matrix that expands each input column with one hot encoding" 212 "based on the bucketization");
214 OPERATOR_SCHEMA(BatchOneHot)
218 Input is a matrix tensor. Its first dimension is the batch 219 size. Expand each column of it using one hot encoding. The `lengths` specifies 220 the size of each column after encoding, and the `values` is the dictionary value 221 of one-hot encoding for each column. For example 223 If data = [[2, 3], [4, 1], [2, 5]], lengths = [2, 3], 224 and values = [2, 4, 1, 3, 5], then 226 output = [[1, 0, 0, 1, 0], [0, 1, 1, 0, 0], [1, 0, 0, 0, 1]] 228 .Input(0, "data",
"input tensor matrix")
229 .Input(1,
"lengths",
"the size is the same as the width of the `data`")
230 .Input(2,
"values",
"one hot encoding dictionary values")
234 "output matrix that expands each input column with one hot encoding")
235 .TensorInferenceFunction(TensorInferenceForBatchOneHot)
236 .CostInferenceFunction(
239 OPERATOR_SCHEMA(OneHot)
243 Given a sequence of indices, one for each example in a batch, returns a matrix 244 where each inner dimension has the size of the index and has 1.0 in the index 245 active in the given example, and 0.0 everywhere else. 247 .Input(0, "indices",
"The active index for each example in the batch.")
251 "Scalar with the size of the index. Must be in CPU context")
252 .Output(0,
"one_hots",
"Matrix of size len(indices) x index_size");
254 OPERATOR_SCHEMA(SegmentOneHot)
258 Given a sequence of indices, segmented by the lengths tensor, returns a matrix 259 that has the elements in each sequence set to 1.0, and 0.0 everywhere else. 261 .Input(0, "lengths",
"Size of each segment.")
262 .Input(1,
"indices",
"Active indices, of size sum(lengths)")
263 .Input(2,
"index_size_tensor",
"Size of the index")
264 .Output(0,
"one_hots",
"Matrix of size len(lengths) x index_size");
266 NO_GRADIENT(BatchOneHot);
268 NO_GRADIENT(SegmentOneHot);
269 NO_GRADIENT(BucketBatchOneHot);
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
std::function< struct Cost(const OperatorDef &, const vector< TensorShape > &)> CostInferenceFunctionType
Registers a function that takes in an OperatorDef and a series of input shapes and returns the total ...