/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators/lengths_pad_op.h (2574B)
#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 LengthsPadOp : public Operator {
public:
USE_OPERATOR_CONTEXT_FUNCTIONS;
template
explicit LengthsPadOp(Args&&... args)
: Operator(std::forward(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>::call(
this, Input(DATA));
}
template
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 total_length = 0;
CPUContext cpuContext;
math::Sum(
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());
auto block_size = data.size_from_dim(1);
auto src_data = data.template data();
auto out_data = output->template mutable_data();
math::Set(
output->numel(), static_cast(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(
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_