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
|
#include "caffe2/operators/lengths_pad_op.h"
namespace caffe2 {
REGISTER_CPU_OPERATOR(LengthsPad, LengthsPadOp<CPUContext>);
OPERATOR_SCHEMA(LengthsPad)
.NumInputs(2)
.NumOutputs(1)
.SetDoc(R"DOC(
Given DATA tensor of rank r >= 1, and LENGTHS tensor of rank 1, pad each
segment in DATA with `value`, so that each segment's length is `target_length`.
If will throw, if there is segment of length larger than `target_length`.
Example:
DATA = [
[2.3, 3.4],
[4.5, 5.7],
[6.8, 7.9],
]
LENGTHS = [0, 1, 1, 1]
and target_length = 2, padding value = -1.0
OUTPUT = [
[-1.0, -1.0],
[-1.0, -1.0],
[2.3, 3.4],
[-1.0, -1.0],
[4.5, 5.7],
[-1.0, -1.0],
[6.8, 7.9],
[-1.0, -1.0],
]
)DOC")
.Input(
0,
"DATA",
"Tensor of rank r >= 1. First dimension must be equal to the size of "
"lengths")
.Input(1, "LENGTHS", "Tensor of int32 lengths of rank 1")
.Output(0, "OUTPUT", "Padded DATA tensor")
.Arg("padding_value", "The value to pad the data")
.Arg("target_length", "The target length of each segment")
.TensorInferenceFunction([](const OperatorDef& def,
const vector<TensorShape>& in) {
vector<TensorShape> out(1);
ArgumentHelper helper(def);
int target_length = helper.GetSingleArgument<int>("target_length", -1);
CAFFE_ENFORCE_GE(target_length, 1);
vector<int> output_dims;
const auto& data_dims = GetDimsVector(in[0]);
const auto& lengths_dims = GetDimsVector(in[1]);
output_dims.push_back(lengths_dims[0] * target_length);
output_dims.insert(
output_dims.end(), data_dims.begin() + 1, data_dims.end());
out[0] = CreateTensorShape(output_dims, in[0].data_type());
return out;
});
NO_GRADIENT(LengthsPad);
} // namespace caffe2
|
