Caffe2 - C++ API
A deep learning, cross platform ML framework
key_split_ops.h
1 
17 #pragma once
18 
19 #include <vector>
20 
21 #include "caffe2/core/context.h"
22 #include "caffe2/core/operator.h"
23 #include "caffe2/utils/math.h"
24 
25 namespace caffe2 {
26 template <typename T, class Context>
27 class KeySplitOp : public Operator<Context> {
28  public:
29  USE_OPERATOR_CONTEXT_FUNCTIONS;
30 
31  KeySplitOp(const OperatorDef& operator_def, Workspace* ws)
32  : Operator<Context>(operator_def, ws),
33  categorical_limit_(
34  OperatorBase::GetSingleArgument<int>("categorical_limit", 0)) {
35  CAFFE_ENFORCE_GT(categorical_limit_, 0);
36  }
37 
38  bool RunOnDevice() override {
39  auto& keys = Input(0);
40  int N = keys.size();
41  const T* keys_data = keys.template data<T>();
42  std::vector<int> counts(categorical_limit_);
43  std::vector<int*> eids(categorical_limit_);
44  for (int k = 0; k < categorical_limit_; k++) {
45  counts[k] = 0;
46  }
47  for (int i = 0; i < N; i++) {
48  int k = keys_data[i];
49  CAFFE_ENFORCE_GT(categorical_limit_, k);
50  CAFFE_ENFORCE_GE(k, 0);
51  counts[k]++;
52  }
53  for (int k = 0; k < categorical_limit_; k++) {
54  auto* eid = Output(k);
55  eid->Resize(counts[k]);
56  eids[k] = eid->template mutable_data<int>();
57  counts[k] = 0;
58  }
59  for (int i = 0; i < N; i++) {
60  int k = keys_data[i];
61  eids[k][counts[k]++] = i;
62  }
63  return true;
64  }
65 
66  private:
67  int categorical_limit_;
68 };
69 } // namespace caffe2
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.