mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
f2f43ad2da
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/10974 Pull Request resolved: https://github.com/pytorch/pytorch/pull/10291 This new operator will do the following: Given a LENGTHS vector and n_splits, output a "split" LENGTHS vector where: 1. Each length in input vector is split into n_splits values (thus output vector should have LENGTHS.size(0) * n_splits elements) 2. The new lengths in output should be evenly split, and if the length is not divisible by n_splits, then order new values in descending order. (e.g. n_splits = 3, length = 5 -> 2 2 1) 3. If n_splits > some element in the array, its split elements will contain 0s. (e.g. n_splits = 3, length = 2 - > 1 1 0) Reviewed By: bddppq, chocjy Differential Revision: D9013119 fbshipit-source-id: 82bf3371ec08c41fc3379177f0007afc142e0d84
75 lines
2.2 KiB
C++
75 lines
2.2 KiB
C++
#ifndef CAFFE2_OPERATORS_LENGTH_SPLIT_OP_H_
|
|
#define CAFFE2_OPERATORS_LENGTH_SPLIT_OP_H_
|
|
|
|
#include "caffe2/core/common_omp.h"
|
|
#include "caffe2/core/context.h"
|
|
#include "caffe2/core/logging.h"
|
|
#include "caffe2/core/operator.h"
|
|
#include "caffe2/utils/math.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
template <class Context>
|
|
class LengthsSplitOp final : public Operator<Context> {
|
|
public:
|
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
|
|
|
LengthsSplitOp(const OperatorDef& operator_def, Workspace* ws)
|
|
: Operator<Context>(operator_def, ws),
|
|
n_split_(OperatorBase::GetSingleArgument<int32_t>("n_split", 0)) {
|
|
if (InputSize() == 1) {
|
|
// If not specified, then must have this argument
|
|
CAFFE_ENFORCE(
|
|
OperatorBase::HasArgument("n_split"),
|
|
"Argument `n_split` is missing and was not specified as input.");
|
|
CAFFE_ENFORCE(
|
|
n_split_ > 0,
|
|
"`n_split` must contain a positive value for defined behavior.");
|
|
}
|
|
}
|
|
~LengthsSplitOp() {}
|
|
|
|
bool RunOnDevice() override {
|
|
const auto& L = Input(0);
|
|
CAFFE_ENFORCE_EQ(L.ndim(), 1, "Input `LENGTHS` should be a 1D vector.");
|
|
|
|
if (InputSize() > 1) {
|
|
// We potentially have n_split specified as inputs as well
|
|
CAFFE_ENFORCE(
|
|
Input(1).ndim() == 1 && Input(1).size() == 1,
|
|
"Input `n_split` should be a vector of size 1.");
|
|
|
|
const auto& input1 = Input(1);
|
|
context_.template CopyItems<Context, CPUContext>(
|
|
input1.meta(), 1, input1.raw_data(), &n_split_);
|
|
}
|
|
|
|
CAFFE_ENFORCE(
|
|
n_split_ > 0,
|
|
"`n_split` must contain a positive value for defined behavior.");
|
|
const auto M = L.size();
|
|
|
|
auto* Y = Output(0);
|
|
Y->Resize(M * n_split_);
|
|
|
|
const int32_t* Ldata = L.template data<int32_t>();
|
|
int32_t* Ydata = Y->template mutable_data<int32_t>();
|
|
|
|
for (int i = 0; i < M; i++) {
|
|
int32_t mod = Ldata[i] % n_split_;
|
|
int32_t res =
|
|
mod != 0 ? math::divUp(Ldata[i], n_split_) : Ldata[i] / n_split_ + 1;
|
|
for (int j = 0; j < n_split_; j++) {
|
|
Ydata[(i * n_split_) + j] = mod-- > 0 ? res : res - 1;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
private:
|
|
int32_t n_split_;
|
|
};
|
|
|
|
} // namespace caffe2
|
|
|
|
#endif // CAFFE2_OPERATORS_LENGTH_SPLIT_OP_H_
|