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
|
#pragma once
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
#include "caffe2/utils/eigen_utils.h"
#include "caffe2/utils/math.h"
namespace caffe2 {
namespace detail {
template <typename T, typename Context>
void VariableLengthSequencePadding(
int N,
int B,
int M,
T* X,
const int32_t* seqLengths,
const T padValue,
Context* /*context*/) {
for (const auto j : c10::irange(B)) {
for (int i = seqLengths[j]; i < N; i++) {
EigenVectorArrayMap<T>(X + B * M * i + M * j, M).setConstant(padValue);
}
}
}
} // namespace detail
template <typename T, typename Context>
class VariableLengthSequencePaddingOp : public Operator<Context> {
public:
template <class... Args>
explicit VariableLengthSequencePaddingOp(Args&&... args)
: Operator<Context>(std::forward<Args>(args)...) {}
USE_OPERATOR_CONTEXT_FUNCTIONS;
bool RunOnDevice() override {
const auto N = Input(INPUT).size(0);
const auto B = Input(INPUT).size(1);
const auto M = Input(INPUT).size(2);
auto X = Output(OUTPUT)->template mutable_data<T>();
auto seqLengths = Input(SEQ_LENGTHS).template data<int32_t>();
detail::VariableLengthSequencePadding<T, Context>(
N, B, M, X, seqLengths, 0, &context_);
return true;
}
protected:
INPUT_TAGS(INPUT, SEQ_LENGTHS);
OUTPUT_TAGS(OUTPUT);
};
} // namespace caffe2
|
