1 #include "caffe2/operators/pack_segments.h" 7 bool PackSegmentsOp<CPUContext>::DoRunWithType() {
9 TensorTypes2<char, int32_t, int64_t, float, std::string>,
10 T>::call(
this, Input(DATA));
14 template <
typename T,
typename Data_T>
15 bool PackSegmentsOp<CPUContext>::DoRunWithType2() {
16 const auto& data = Input(DATA);
17 const auto& lengths = Input(LENGTHS);
18 auto* output = Output(0);
20 if (return_presence_mask_) {
21 presence_mask = Output(1);
24 CAFFE_ENFORCE(data.ndim() >= 1,
"DATA should be at least 1-D");
25 CAFFE_ENFORCE(lengths.ndim() == 1,
"LENGTH should be 1-D");
28 const T* l = lengths.template data<T>();
30 TIndex total_length = 0;
31 for (T i = 0; i < lengths.dim(0); ++i) {
32 max_length = std::max(max_length, l[i]);
40 " PackSegments requires that the sum of the lengths ",
42 " is equal to the first data dimension ",
45 auto shape = data.dims();
46 shape[0] = max_length;
47 shape.insert(shape.begin(), lengths.size());
48 output->Resize(shape);
51 auto* out =
static_cast<char*
>(output->raw_mutable_data(data.meta()));
53 bool* presence_mask_data =
nullptr;
54 if (return_presence_mask_) {
56 std::vector<caffe2::TIndex> presence_shape{lengths.size(), max_length};
57 presence_mask->Resize(presence_shape);
58 presence_mask_data = presence_mask->template mutable_data<bool>();
67 if (output->template IsType<float>()) {
68 math::Set<float, CPUContext>(
71 output->template mutable_data<float>(),
74 if (return_presence_mask_) {
75 memset(presence_mask_data, (
int)
false, presence_mask->size());
78 auto block_size = data.size_from_dim(1);
79 auto block_bytesize = data.itemsize() * block_size;
80 const auto* d =
static_cast<const char*
>(data.raw_data());
82 for (TIndex i = 0; i < lengths.dim(0); ++i) {
83 context_.template CopyItems<CPUContext, CPUContext>(
86 d + block_bytesize * start,
87 out + block_bytesize * max_length * i);
88 if (return_presence_mask_) {
89 memset(presence_mask_data + max_length * i, (
int)
true, l[i]);
99 bool UnpackSegmentsOp<CPUContext>::DoRunWithType() {
100 return DispatchHelper<
101 TensorTypes2<char, int32_t, int64_t, float, std::string>,
102 T>::call(
this, Input(DATA));
106 template <
typename T,
typename Data_T>
107 bool UnpackSegmentsOp<CPUContext>::DoRunWithType2() {
108 const auto& data = Input(DATA);
109 const auto& lengths = Input(LENGTHS);
110 auto* output = Output(0);
112 CAFFE_ENFORCE(data.ndim() >= 2,
"DATA should be at least 2-D");
113 CAFFE_ENFORCE(lengths.ndim() == 1,
"LENGTH should be 1-D");
115 const T* l = lengths.template data<T>();
117 TIndex total_l = std::accumulate(l, l + lengths.dim(0), (TIndex)0);
119 auto shape = data.dims();
121 shape[0] == lengths.dim(0),
"LENGTH should match DATA in dimension 0");
122 shape.erase(shape.begin());
124 output->Resize(shape);
126 auto* out =
static_cast<char*
>(output->raw_mutable_data(data.meta()));
127 if (!(data.dim(0) * data.dim(1))) {
130 auto block_size = data.size_from_dim(2);
131 auto block_bytesize = data.itemsize() * block_size;
132 const auto* d =
static_cast<const char*
>(data.raw_data());
134 for (TIndex i = 0; i < lengths.dim(0); ++i) {
135 context_.template CopyItems<CPUContext, CPUContext>(
138 d + block_bytesize * data.dim(1) * i,
139 out + block_bytesize * start);
145 REGISTER_CPU_OPERATOR(PackSegments, PackSegmentsOp<CPUContext>);
146 REGISTER_CPU_OPERATOR(UnpackSegments, UnpackSegmentsOp<CPUContext>);
148 OPERATOR_SCHEMA(PackSegments)
152 "Map N dim tensor to N+1 dim based on length blob. Sequences that \ 153 are shorter than the longest sequence are padded with zeros.")
157 "1-d int/long tensor contains the length in each of the output.")
158 .Input(1,
"tensor",
"N dim Tensor.")
163 "where dim(1) is the max length" 164 ", dim(0) is the batch size.")
168 "2 dim boolean tensor" 169 ", false where packed_tensor is padded, true otherwise.")
172 "Padding number in the packed segments. Use true to pad \ 173 -infinity, otherwise pad zeros")
175 "return_presence_mask",
176 "bool whether to return presence mask, false by default");
177 OPERATOR_SCHEMA(UnpackSegments)
180 .SetDoc(
"Map N+1 dim tensor to N dim based on length blob")
184 "1-d int/long tensor contains the length in each of the input.")
185 .Input(1,
"tensor",
"N+1 dim Tensor.")
186 .Output(0,
"packed_tensor",
"N dim Tensor");
189 using GradientMakerBase::GradientMakerBase;
190 vector<OperatorDef> GetGradientDefs()
override {
194 vector<string>{I(0), GO(0)},
195 vector<string>{GI(1)});
201 using GradientMakerBase::GradientMakerBase;
202 vector<OperatorDef> GetGradientDefs()
override {
204 "PackSegments",
"", vector<string>{I(0), GO(0)}, vector<string>{GI(1)});
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
static vector< OperatorDef > SingleGradientDef(const Args &...args)
a helper function to allow one to create one single operator def, which is usually the case for many ...