#ifndef CAFFE2_OPERATORS_LENGTHS_PAD_OP_H_
|
#define CAFFE2_OPERATORS_LENGTHS_PAD_OP_H_
|
|
#include "caffe2/core/operator.h"
|
#include "caffe2/utils/math.h"
|
|
namespace caffe2 {
|
|
template <class Context>
|
class LengthsPadOp : public Operator<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
template <class... Args>
|
explicit LengthsPadOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
OP_SINGLE_ARG(double, "padding_value", padding_value_, -1),
|
OP_SINGLE_ARG(int, "target_length", target_length_, -1) {
|
CAFFE_ENFORCE_GE(target_length_, 1, "target_length argument must be >= 1");
|
}
|
|
bool RunOnDevice() override {
|
return DispatchHelper<TensorTypes<float, double, int32_t, int64_t>>::call(
|
this, Input(DATA));
|
}
|
|
template <typename T>
|
bool DoRunWithType() {
|
auto& data = Input(DATA);
|
auto& lengths = Input(LENGTHS);
|
|
CAFFE_ENFORCE_EQ(lengths.dim(), 1, "LENGTHS must be 1-D");
|
CAFFE_ENFORCE_GE(data.dim(), 1, "DATA should be at least 1-D");
|
|
// Context::CopyFrom and math::Sum need the same context to avoid race
|
// conditions
|
// why? CPUContext is not used in Sum
|
lengths_host_.CopyFrom(lengths);
|
|
auto lengths_size = lengths_host_.numel();
|
auto* lengths_data = lengths_host_.template data<int32_t>();
|
|
int32_t total_length = 0;
|
CPUContext cpuContext;
|
math::Sum<int32_t, CPUContext>(
|
lengths_size, lengths_data, &total_length, &cpuContext);
|
|
CAFFE_ENFORCE_EQ(total_length, data.size(0));
|
|
auto shape = data.sizes().vec();
|
shape[0] = lengths_size * target_length_;
|
auto* output = Output(0, shape, at::dtype<T>());
|
|
auto block_size = data.size_from_dim(1);
|
auto src_data = data.template data<T>();
|
auto out_data = output->template mutable_data<T>();
|
|
math::Set(
|
output->numel(), static_cast<T>(padding_value_), out_data, &context_);
|
for (int64_t i = 0; i < lengths_size; ++i) {
|
auto length = lengths_data[i];
|
CAFFE_ENFORCE_GE(length, 0);
|
CAFFE_ENFORCE_GE(
|
target_length_,
|
length,
|
"Length at index = ",
|
i,
|
" is larger than target length");
|
|
context_.template CopySameDevice<T>(
|
block_size * length, src_data, out_data);
|
|
out_data += block_size * target_length_;
|
src_data += block_size * length;
|
}
|
return true;
|
}
|
|
INPUT_TAGS(DATA, LENGTHS);
|
|
private:
|
double padding_value_;
|
int target_length_;
|
Tensor lengths_host_{CPU};
|
};
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_OPERATORS_LENGTHS_PAD_OP_H_
|
