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
|
#include "caffe2/operators/sinusoid_position_encoding_op.h"
namespace caffe2 {
REGISTER_CPU_OPERATOR(
SinusoidPositionEncoding,
SinusoidPositionEncodingOp<CPUContext>);
OPERATOR_SCHEMA(SinusoidPositionEncoding)
.NumInputs(1)
.NumOutputs(1)
.SetDoc(R"DOC(
Calculates a sinusoid position encoding tensor as described
in https://arxiv.org/abs/1706.03762. Takes a 2-D tensor
(of size M x K) of positions as input, the embedding size
as an argument, and outputs a position encoding tensor of
size (M x K x embedding_size). Here M is typically the max
sequence length and K is typically the batch size.
The input tensor must satisfy input[m, 0] == input[m, k] for all k.
Encoded as amplitude * SIN(pos/alpha^(i/embedding_size)) if i is even,
else amplitude * COS(pos/alpha^(i/embedding_size)). Here, pos is the position,
alpha and amplitude are tuning parameters, i is the current dimension for
the embedding, and embedding_size is the number of total dimensions in
the embedding.
)DOC")
.Arg(
"embedding_size",
"Desired embedding size/number of dimensions -- defaults to 100")
.Arg("alpha", "Sinusoid tuning parameter -- defaults to 10000")
.Arg("amplitude", "Amplitude of Sin/Cos output")
.Input(0, "positions", "2-D tensor of positions to be encoded")
.Output(0, "output", "3-D tensor representing the positional encoding");
} // namespace caffe2
|
