#ifndef CAFFE2_OPERATORS_SPACE_BATCH_OP_H_
|
#define CAFFE2_OPERATORS_SPACE_BATCH_OP_H_
|
|
#include "caffe2/core/context.h"
|
#include "caffe2/core/logging.h"
|
#include "caffe2/core/operator.h"
|
#include "caffe2/utils/math.h"
|
|
namespace caffe2 {
|
|
template <typename Context>
|
void spaceToBatch(
|
const Tensor& input,
|
int pad_t,
|
int pad_l,
|
int block_size,
|
Tensor* output,
|
Context* /*context*/) {
|
CAFFE_ENFORCE(input.dim() == 4);
|
CAFFE_ENFORCE(output->dim() == 4);
|
|
const int output_batch = output->dim32(0);
|
const int output_depth = output->dim32(1);
|
const int output_height = output->dim32(2);
|
const int output_width = output->dim32(3);
|
|
const int input_batch = input.dim32(0);
|
const int input_depth = input.dim32(1);
|
const int input_height = input.dim32(2);
|
const int input_width = input.dim32(3);
|
|
for (int out_b = 0; out_b < output_batch; ++out_b) {
|
const int in_b = out_b % input_batch;
|
const int offset_w = (out_b / input_batch) % block_size;
|
const int offset_h = (out_b / input_batch) / block_size;
|
for (int d = 0; d < input_depth; ++d) {
|
for (int out_h = 0; out_h < output_height; ++out_h) {
|
const int in_h = out_h * block_size + offset_h - pad_t;
|
for (int out_w = 0; out_w < output_width; ++out_w) {
|
const int in_w = out_w * block_size + offset_w - pad_l;
|
const auto output_offset =
|
((out_b * output_depth + d) * output_height + out_h) *
|
output_width +
|
out_w;
|
const auto input_offset =
|
((in_b * input_depth + d) * input_height + in_h) * input_width +
|
in_w;
|
if (in_h >= 0 && in_w >= 0 && in_h < input_height &&
|
in_w < input_width) {
|
output->template mutable_data<float>()[output_offset] =
|
input.template data<float>()[input_offset];
|
} else {
|
output->template mutable_data<float>()[output_offset] = 0.0;
|
}
|
}
|
}
|
}
|
}
|
}
|
|
template <typename Context>
|
void batchToSpace(
|
const Tensor& input,
|
int pad_t,
|
int pad_l,
|
int block_size,
|
Tensor* output,
|
Context* /*context*/) {
|
CAFFE_ENFORCE(input.dim() == 4);
|
CAFFE_ENFORCE(output->dim() == 4);
|
|
const int output_batch = output->dim32(0);
|
const int output_depth = output->dim32(1);
|
const int output_height = output->dim32(2);
|
const int output_width = output->dim32(3);
|
|
const int input_batch = input.dim32(0);
|
const int input_depth = input.dim32(1);
|
const int input_height = input.dim32(2);
|
const int input_width = input.dim32(3);
|
|
CAFFE_ENFORCE(input_depth == output_depth);
|
for (int in_b = 0; in_b < input_batch; ++in_b) {
|
const int out_b = in_b % output_batch;
|
const int offset_w = (in_b / output_batch) % block_size;
|
const int offset_h = (in_b / output_batch) / block_size;
|
for (int d = 0; d < input_depth; ++d) {
|
for (int in_h = 0; in_h < input_height; ++in_h) {
|
const int out_h = in_h * block_size + offset_h - pad_t;
|
for (int in_w = 0; in_w < input_width; ++in_w) {
|
const int out_w = in_w * block_size + offset_w - pad_l;
|
if (out_h >= 0 && out_w >= 0 && out_h < output_height &&
|
out_w < output_width) {
|
const auto output_offset =
|
((out_b * output_depth + d) * output_height + out_h) *
|
output_width +
|
out_w;
|
const auto input_offset =
|
((in_b * input_depth + d) * input_height + in_h) * input_width +
|
in_w;
|
output->template mutable_data<float>()[output_offset] =
|
input.template data<float>()[input_offset];
|
}
|
}
|
}
|
}
|
}
|
}
|
|
template <typename Context>
|
class SpaceBatchOpBase : public Operator<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
template <class... Args>
|
explicit SpaceBatchOpBase(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
pad_(this->template GetSingleArgument<int>("pad", 0)),
|
pad_t_(this->template GetSingleArgument<int>("pad_t", pad_)),
|
pad_l_(this->template GetSingleArgument<int>("pad", pad_)),
|
pad_b_(this->template GetSingleArgument<int>("pad", pad_)),
|
pad_r_(this->template GetSingleArgument<int>("pad", pad_)),
|
block_size_(this->template GetSingleArgument<int>("block_size", 2)),
|
order_(StringToStorageOrder(
|
this->template GetSingleArgument<string>("order", "NCHW"))) {
|
CAFFE_ENFORCE(order_ == StorageOrder::NCHW);
|
}
|
|
protected:
|
int pad_;
|
int pad_t_;
|
int pad_l_;
|
int pad_b_;
|
int pad_r_;
|
int block_size_;
|
StorageOrder order_;
|
};
|
|
template <typename Context>
|
class SpaceToBatchOp final : public SpaceBatchOpBase<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
using SpaceBatchOpBase<Context>::SpaceBatchOpBase;
|
|
bool RunOnDevice() override {
|
const auto& input = Input(0);
|
auto* output = Output(0);
|
const int batch = input.dim32(0);
|
const int depth = input.dim32(1);
|
const int height = this->pad_b_ + this->pad_t_ + input.dim32(2);
|
const int width = this->pad_l_ + this->pad_r_ + input.dim32(3);
|
CAFFE_ENFORCE(
|
height % this->block_size_ == 0,
|
"Height: ",
|
height,
|
", block size: ",
|
this->block_size_);
|
CAFFE_ENFORCE(width % this->block_size_ == 0);
|
|
const int output_batch = batch * this->block_size_ * this->block_size_;
|
const int output_height = height / this->block_size_;
|
const int output_width = width / this->block_size_;
|
Output(0)->Resize(output_batch, depth, output_height, output_width);
|
|
spaceToBatch<Context>(
|
input,
|
this->pad_t_,
|
this->pad_l_,
|
this->block_size_,
|
output,
|
&context_);
|
|
return true;
|
}
|
};
|
|
template <typename Context>
|
class BatchToSpaceOp final : public SpaceBatchOpBase<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
using SpaceBatchOpBase<Context>::SpaceBatchOpBase;
|
|
bool RunOnDevice() override {
|
const auto& input = Input(0);
|
auto* output = Output(0);
|
const int batch = input.dim32(0);
|
const int depth = input.dim32(1);
|
const int height = input.dim32(2);
|
const int width = input.dim32(3);
|
|
const int output_batch = batch / this->block_size_ / this->block_size_;
|
const int output_height =
|
height * this->block_size_ - this->pad_b_ - this->pad_t_;
|
const int output_width =
|
width * this->block_size_ - this->pad_l_ - this->pad_r_;
|
Output(0)->Resize(output_batch, depth, output_height, output_width);
|
batchToSpace<Context>(
|
input,
|
this->pad_t_,
|
this->pad_l_,
|
this->block_size_,
|
output,
|
&context_);
|
return true;
|
}
|
};
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_OPERATORS_SPACE_BATCH_OP_H_
|
