1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
|
#ifndef CAFFE2_OPERATORS_SINUSOID_POSITION_ENCODING_OP_H_
#define CAFFE2_OPERATORS_SINUSOID_POSITION_ENCODING_OP_H_
#ifdef _MSC_VER
#ifndef _USE_MATH_DEFINES
#define _USE_MATH_DEFINES
#endif
#endif // _MSC_VER
#include <cmath>
#include "caffe2/core/operator.h"
#include "Eigen/Core"
#include "caffe2/utils/eigen_utils.h"
namespace caffe2 {
template <class Context>
class SinusoidPositionEncodingOp : public Operator<Context> {
public:
template <class... Args>
explicit SinusoidPositionEncodingOp(Args&&... args)
: Operator<Context>(std::forward<Args>(args)...),
embedding_size_(
this->template GetSingleArgument<int>("embedding_size", 100)),
alpha_(this->template GetSingleArgument<float>("alpha", 10000)),
amplitude_(this->template GetSingleArgument<float>("amplitude", 1)) {}
USE_OPERATOR_CONTEXT_FUNCTIONS;
bool RunOnDevice() override {
return DispatchHelper<TensorTypes<int32_t, int64_t>>::call(
this, this->template Input<Tensor>(0, CPU));
}
template <typename Index>
bool DoRunWithType() {
auto& positions = Input(0);
CAFFE_ENFORCE_EQ(positions.dim(), 2, "POSITIONS should be a 2-D tensor");
auto shape = positions.sizes().vec();
shape.push_back(embedding_size_);
auto* output = Output(0, shape, at::dtype<float>());
int M = shape[0];
int K = shape[1];
const Index* idxs = positions.template data<Index>();
float* out = output->template mutable_data<float>();
float log_alpha = std::log(alpha_);
float max_alpha_pow =
((float)embedding_size_ - 1.0f) / (float)embedding_size_;
for (const auto i : c10::irange(M)) {
float pos = (float)idxs[i * K];
// Compute the embedding for position i, example 0 first
float* row = &out[i * K * embedding_size_];
Eigen::Map<Eigen::VectorXf> row_map(row, embedding_size_, 1);
auto row_array = row_map.array();
float log_pos = std::log(pos);
row_array.setLinSpaced(
embedding_size_, log_pos, log_pos - log_alpha * max_alpha_pow);
row_array = row_array.exp().eval();
// row_array[k] == pos / alpha^(k / embedding_size)
// Phase shift so that alternating elements are cosines
for (int k = 1; k < embedding_size_; k += 2) {
row[k] += (float)M_PI_2;
}
row_array = amplitude_ * row_array.sin().eval();
// Copy the embedding to position i in the other examples
for (const auto j : c10::irange(1, K)) {
int base = i * K * embedding_size_;
std::copy(
&out[base],
&out[base + embedding_size_],
&out[base + j * embedding_size_]);
}
}
return true;
}
protected:
int embedding_size_;
float alpha_;
float amplitude_;
};
} // namespace caffe2
#endif // CAFFE2_OPERATORS_SINUSOID_POSITION_ENCODING_OP_H_
|
