#ifndef CAFFE2_OPERATORS_REDUCE_FRONT_BACK_SUM_MEAN_OPS_H_
|
#define CAFFE2_OPERATORS_REDUCE_FRONT_BACK_SUM_MEAN_OPS_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, bool FIRSTDIMS, bool NORMALIZE>
|
class SumReduceDimsOp final : public Operator<Context> {
|
public:
|
template <class... Args>
|
explicit SumReduceDimsOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
num_reduce_dims_(
|
this->template GetSingleArgument<int32_t>("num_reduce_dim", 1)) {}
|
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
|
bool RunOnDevice() override {
|
return DispatchHelper<TensorTypes<int, int64_t, float, double>>::call(
|
this, Input(0));
|
}
|
|
template <typename T>
|
bool DoRunWithType() {
|
auto& X = Input(0);
|
|
CAFFE_ENFORCE(
|
num_reduce_dims_ >= 0 && num_reduce_dims_ <= X.dim(),
|
"For N-dim input tensor, support num_reduce_dims in range [0, N].");
|
|
vector<int64_t> output_shape;
|
int start_index = FIRSTDIMS ? num_reduce_dims_ : 0;
|
int end_index =
|
FIRSTDIMS ? X.dim() : X.dim() - num_reduce_dims_;
|
for (int i = start_index; i < end_index; ++i) {
|
output_shape.push_back(X.sizes()[i]);
|
}
|
auto* Y = Output(0, output_shape, at::dtype<T>());
|
|
const int rows = FIRSTDIMS ? X.size_to_dim(num_reduce_dims_)
|
: X.size_to_dim(X.dim() - num_reduce_dims_);
|
const int cols = FIRSTDIMS ? X.size_from_dim(num_reduce_dims_)
|
: X.size_from_dim(X.dim() - num_reduce_dims_);
|
|
const T* in_data = X.template data<T>();
|
T* out_data = Y->template mutable_data<T>();
|
|
if (cols == 0 || rows == 0) {
|
math::Set(Y->numel(), static_cast<T>(0), out_data, &context_);
|
return true;
|
}
|
|
const int32_t* lengths_data = nullptr;
|
if (InputSize() > 1) {
|
const auto& lengths = Input(1);
|
lengths_data = lengths.template data<int32_t>();
|
CAFFE_ENFORCE(
|
num_reduce_dims_ == 1,
|
"Given lengths input, the number of reduce dimensions should be one.");
|
const int batch_size = FIRSTDIMS ? cols : rows;
|
CAFFE_ENFORCE(
|
lengths.numel() == batch_size,
|
"The size of lengths vector doesn't match the batch size.");
|
}
|
|
Compute(rows, cols, in_data, lengths_data, out_data);
|
|
return true;
|
}
|
|
private:
|
template <typename T>
|
void Compute(
|
int rows,
|
int cols,
|
const T* in_data,
|
const int32_t* lengths_data,
|
T* out_data);
|
|
int num_reduce_dims_;
|
};
|
|
template <class Context, bool FIRSTDIMS, bool NORMALIZE>
|
class SumReduceDimsGradientOp final : public Operator<Context> {
|
public:
|
template <class... Args>
|
explicit SumReduceDimsGradientOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
num_reduce_dims_(
|
this->template GetSingleArgument<int32_t>("num_reduce_dim", 1)) {}
|
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
|
bool RunOnDevice() override {
|
return DispatchHelper<TensorTypes<int, long, float, double>>::call(
|
this, Input(0));
|
}
|
|
template <typename T>
|
bool DoRunWithType() {
|
auto& dY = Input(0);
|
auto& input_1 = Input(1);
|
|
vector<int64_t> dX_sizes;
|
// In previous diff we changed the semantic: Input(1) was changed from
|
// the shape of the input to the data tensor. This made the backward
|
// computation incompatible with old models. To fix this, we check
|
// the dimension and type of Input(1).
|
if (input_1.dim() == 1 && input_1.template IsType<int64_t>()) {
|
// Input(1) is the shape of the input
|
shape_.CopyFrom(input_1);
|
// Copy first dims
|
dX_sizes = vector<int64_t>(
|
shape_.template data<int64_t>(),
|
shape_.template data<int64_t>() + shape_.numel());
|
} else {
|
// Input(1) is data tensor X
|
dX_sizes = input_1.sizes().vec();
|
}
|
auto* dX = Output(0, dX_sizes, at::dtype<T>());
|
|
const int rows = FIRSTDIMS ? dX->size_to_dim(num_reduce_dims_)
|
: dX->size_to_dim(dX->dim() - num_reduce_dims_);
|
const int cols = FIRSTDIMS
|
? dX->size_from_dim(num_reduce_dims_)
|
: dX->size_from_dim(dX->dim() - num_reduce_dims_);
|
|
const int32_t* lengths_data = nullptr;
|
if (InputSize() > 2) {
|
const auto& lengths = Input(2);
|
lengths_data = lengths.template data<int32_t>();
|
CAFFE_ENFORCE(
|
num_reduce_dims_ == 1,
|
"Given lengths input, the number of reduce dimensions should be one.");
|
const int batch_size = FIRSTDIMS ? cols : rows;
|
CAFFE_ENFORCE(
|
lengths.numel() == batch_size,
|
"The size of lengths vector doesn't match the batch size.");
|
}
|
|
const T* dYdata = dY.template data<T>();
|
T* dXdata = dX->template mutable_data<T>();
|
Compute<T>(rows, cols, dYdata, lengths_data, dXdata);
|
return true;
|
}
|
|
private:
|
template <typename T>
|
void Compute(
|
int rows,
|
int cols,
|
const T* dYdata,
|
const int32_t* lengths_data,
|
T* dXdata);
|
int num_reduce_dims_;
|
// scratch space used for former version of this reducer
|
Tensor shape_{Context::GetDeviceType()};
|
};
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_OPERATORS_REDUCE_FRONT_BACK_SUM_MEAN_OPS_H_
|
