#ifndef CAFFE2_OPERATORS_RESHAPE_OP_H_
|
#define CAFFE2_OPERATORS_RESHAPE_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 {
|
|
// Takes a shape and data tensor and reshapes it
|
template <typename F, class Context>
|
class ReshapeOp : public Operator<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
template <class... Args>
|
explicit ReshapeOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
new_shape_(this->template GetRepeatedArgument<int64_t>("shape")) {}
|
|
bool RunOnDevice() override {
|
if (InputSize() == 2) {
|
return DispatchHelper<TensorTypes<int, int64_t>>::call(this, Input(1));
|
}
|
CAFFE_ENFORCE(
|
OperatorBase::HasArgument("shape"), "Argument `shape` is missing.");
|
return this->template DoRunWithType<int64_t>();
|
}
|
|
template <typename T>
|
bool DoRunWithType() {
|
DoRunWithTypeImpl<T>(Input(0), Output(0));
|
return true;
|
}
|
|
protected:
|
template <typename T>
|
void DoRunWithTypeImpl(const Tensor& input, Tensor* output) {
|
vector<int64_t> actual_new_shape = new_shape_;
|
if (InputSize() == 2) {
|
CAFFE_ENFORCE(
|
!OperatorBase::HasArgument("shape"),
|
"New shape is specified by the input blob, do not pass in "
|
"the argument `shape`.");
|
|
auto& shape = Input(1);
|
CAFFE_ENFORCE(shape.dim() == 1, "Shape should be 1-D");
|
|
const T* shape_data = shape.template data<T>();
|
|
// Bit awkward, but needed so works on both CPU and CUDA contexts
|
std::vector<T> tmpv(shape.numel());
|
if (shape.numel() > 0) {
|
context_.CopyBytesToCPU(
|
shape.numel() * sizeof(T), shape_data, &tmpv[0]);
|
actual_new_shape.assign(tmpv.begin(), tmpv.begin() + shape.numel());
|
}
|
}
|
|
// Checks if the new shape is valid and fills in the missing dimension
|
// specified by -1.
|
// NOTE: At most one dimension can be -1.
|
auto total_size = input.numel();
|
T size = 1;
|
|
// NOTE: support for legacy caffe1 syntax
|
// Copy over the dimensions for those that are specified zero.
|
if (total_size != 0) {
|
for (size_t i = 0; i < actual_new_shape.size() && i < input.dim(); ++i) {
|
if (actual_new_shape[i] == 0) {
|
actual_new_shape[i] = input.size(i);
|
}
|
}
|
}
|
|
int unknown_idx = -1;
|
for (int i = 0; i < actual_new_shape.size(); ++i) {
|
const auto dim = actual_new_shape[i];
|
if (dim == -1) {
|
CAFFE_ENFORCE(
|
unknown_idx == -1,
|
"Argument `shape` has more than one missing dimension.");
|
unknown_idx = i;
|
} else {
|
size *= dim;
|
}
|
}
|
if (size == 0 && total_size != 0) {
|
CAFFE_THROW(
|
"Can not reshape a non-zero size (",
|
total_size,
|
") tensor to zero size.");
|
}
|
if (total_size != 0) {
|
// if tensor is not empty, infer the size of the unknown index
|
if (unknown_idx != -1) {
|
CAFFE_ENFORCE_NE(
|
size,
|
0,
|
"New shape at dim ",
|
unknown_idx,
|
" can not be inferred since new size is zero.");
|
CAFFE_ENFORCE(
|
total_size % size == 0,
|
"Argument `shape` does not agree with the input data.",
|
" (",
|
total_size,
|
" vs ",
|
size,
|
")");
|
actual_new_shape[unknown_idx] = total_size / size;
|
} else {
|
CAFFE_ENFORCE_EQ(
|
total_size,
|
size,
|
"Argument `shape` does not agree with the input data.",
|
" (",
|
total_size,
|
" != ",
|
size,
|
")");
|
}
|
} else if (unknown_idx != -1) {
|
// if size is empty, then set unknown index to be 0 (empty tensor)
|
actual_new_shape[unknown_idx] = 0;
|
}
|
|
// Write the original shape to the second output.
|
auto* old_shape = Output(1, {input.dim()}, at::dtype<T>());
|
T* old_shape_data = old_shape->template mutable_data<T>();
|
for (int i = 0; i < input.dim(); ++i) {
|
math::Set<T, Context>(1, input.size(i), old_shape_data + i, &context_);
|
}
|
|
output->Resize(actual_new_shape);
|
if (output != &input) {
|
// If we are not doing in-place computation, a copy is needed.
|
context_.CopyItemsSameDevice(
|
input.dtype(),
|
input.numel(),
|
input.raw_data(),
|
output->raw_mutable_data(input.dtype()));
|
}
|
}
|
|
private:
|
vector<int64_t> new_shape_;
|
};
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_OPERATORS_RESHAPE_OP_H_
|
