Caffe2 - C++ API
A deep learning, cross platform ML framework
sinusoid_position_encoding_op.h
1 
17 #ifndef CAFFE2_OPERATORS_SINUSOID_POSITION_ENCODING_OP_H_
18 #define CAFFE2_OPERATORS_SINUSOID_POSITION_ENCODING_OP_H_
19 
20 #ifdef _MSC_VER
21 #define _USE_MATH_DEFINES
22 #endif // _MSC_VER
23 #include <cmath>
24 
25 #include "caffe2/core/operator.h"
26 
27 #include "Eigen/Core"
28 
29 namespace caffe2 {
30 
31 template <class Context>
32 class SinusoidPositionEncodingOp : public Operator<Context> {
33  public:
34  SinusoidPositionEncodingOp(const OperatorDef& operator_def, Workspace* ws)
35  : Operator<Context>(operator_def, ws),
36  embedding_size_(OperatorBase::template GetSingleArgument<int>(
37  "embedding_size",
38  100)),
39  alpha_(OperatorBase::template GetSingleArgument<float>("alpha", 10000)),
40  amplitude_(
41  OperatorBase::template GetSingleArgument<float>("amplitude", 1)) {}
42  USE_OPERATOR_CONTEXT_FUNCTIONS;
43 
44  bool RunOnDevice() override {
46  this, OperatorBase::Input<TensorCPU>(0));
47  }
48 
49  template <typename Index>
50  bool DoRunWithType() {
51  auto& positions = Input(0);
52  auto* output = Output(0);
53 
54  CAFFE_ENFORCE_EQ(positions.ndim(), 2, "POSITIONS should be a 2-D tensor");
55 
56  auto shape = positions.dims();
57  shape.push_back(embedding_size_);
58  output->Resize(shape);
59 
60  int M = shape[0];
61  int K = shape[1];
62  const Index* idxs = positions.template data<Index>();
63  float* out = output->template mutable_data<float>();
64 
65  float log_alpha = std::log(alpha_);
66  float max_alpha_pow =
67  ((float)embedding_size_ - 1.0f) / (float)embedding_size_;
68 
69  for (int i = 0; i < M; ++i) {
70  float pos = (float)idxs[i * K];
71 
72  // Compute the embedding for position i, example 0 first
73  float* row = &out[i * K * embedding_size_];
74  Eigen::Map<Eigen::VectorXf> row_map(row, embedding_size_, 1);
75  auto row_array = row_map.array();
76 
77  float log_pos = std::log(pos);
78  row_array.setLinSpaced(
79  embedding_size_, log_pos, log_pos - log_alpha * max_alpha_pow);
80  row_array = row_array.exp().eval();
81  // row_array[k] == pos / alpha^(k / embedding_size)
82 
83  // Phase shift so that alternating elements are cosines
84  for (int k = 1; k < embedding_size_; k += 2) {
85  row[k] += (float)M_PI_2;
86  }
87  row_array = amplitude_ * row_array.sin().eval();
88 
89  // Copy the embedding to position i in the other examples
90  for (int j = 1; j < K; ++j) {
91  int base = i * K * embedding_size_;
92  std::copy(
93  &out[base],
94  &out[base + embedding_size_],
95  &out[base + j * embedding_size_]);
96  }
97  }
98  return true;
99  }
100 
101  protected:
102  int embedding_size_;
103  float alpha_;
104  float amplitude_;
105 };
106 
107 } // namespace caffe2
108 
109 #endif // CAFFE2_OPERATORS_SINUSOID_POSITION_ENCODING_OP_H_
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.