#ifndef CAFFE2_OPERATORS_IM2COL_OP_H_
|
#define CAFFE2_OPERATORS_IM2COL_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 T, class Context>
|
class Im2ColOp final : public Operator<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
template <class... Args>
|
explicit Im2ColOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
pad_(this->template GetSingleArgument<int>("pad", 0)),
|
kernel_h_(this->template GetSingleArgument<int>(
|
"kernel_h",
|
this->template GetSingleArgument<int>("kernel", 0))),
|
kernel_w_(this->template GetSingleArgument<int>(
|
"kernel_w",
|
this->template GetSingleArgument<int>("kernel", 0))),
|
dilation_h_(this->template GetSingleArgument<int>(
|
"dilation_h",
|
this->template GetSingleArgument<int>("dilation", 1))),
|
dilation_w_(this->template GetSingleArgument<int>(
|
"dilation_w",
|
this->template GetSingleArgument<int>("dilation", 1))),
|
stride_h_(this->template GetSingleArgument<int>(
|
"stride_h",
|
this->template GetSingleArgument<int>("stride", 1))),
|
stride_w_(this->template GetSingleArgument<int>(
|
"stride_w",
|
this->template GetSingleArgument<int>("stride", 1))),
|
order_(StringToStorageOrder(
|
this->template GetSingleArgument<string>("order", "NCHW"))) {
|
CAFFE_ENFORCE(kernel_h_ > 0);
|
CAFFE_ENFORCE(kernel_w_ > 0);
|
CAFFE_ENFORCE(dilation_h_ > 0);
|
CAFFE_ENFORCE(dilation_w_ > 0);
|
CAFFE_ENFORCE(stride_h_ > 0);
|
CAFFE_ENFORCE(stride_w_ > 0);
|
CAFFE_ENFORCE(pad_ >= 0);
|
}
|
|
bool RunOnDevice() override {
|
auto& X = Input(0);
|
|
CAFFE_ENFORCE(4 == X.dim());
|
|
int N = 0, C = 0, H = 0, W = 0;
|
switch (order_) {
|
case StorageOrder::NCHW:
|
N = X.dim32(0);
|
C = X.dim32(1);
|
H = X.dim32(2);
|
W = X.dim32(3);
|
break;
|
case StorageOrder::NHWC:
|
N = X.dim32(0);
|
H = X.dim32(1);
|
W = X.dim32(2);
|
C = X.dim32(3);
|
break;
|
default:
|
CAFFE_THROW("Unknown storage order: ", order_);
|
}
|
|
const int dkernel_h = dilation_h_ * (kernel_h_ - 1) + 1;
|
const int dkernel_w = dilation_w_ * (kernel_w_ - 1) + 1;
|
CAFFE_ENFORCE(H >= dkernel_h);
|
CAFFE_ENFORCE(W >= dkernel_w);
|
const int out_h = (H + 2 * pad_ - dkernel_h) / stride_h_ + 1;
|
const int out_w = (W + 2 * pad_ - dkernel_w) / stride_w_ + 1;
|
|
switch (order_) {
|
case StorageOrder::NCHW: {
|
auto* Y = Output(
|
0,
|
std::vector<int64_t>{N, C * kernel_h_ * kernel_w_, out_h, out_w},
|
at::dtype<T>());
|
|
const size_t dx = X.numel() / N;
|
const size_t dy = Y->numel() / N;
|
for (int n = 0; n < N; ++n) {
|
const auto* xdata = X.template data<T>() + (n * dx);
|
auto* ydata = Y->template mutable_data<T>() + (n * dy);
|
math::Im2Col<T, Context, StorageOrder::NCHW>(
|
C,
|
H,
|
W,
|
kernel_h_,
|
kernel_w_,
|
dilation_h_,
|
dilation_w_,
|
pad_,
|
pad_,
|
pad_,
|
pad_,
|
stride_h_,
|
stride_w_,
|
xdata,
|
ydata,
|
&context_);
|
}
|
}; break;
|
case StorageOrder::NHWC: {
|
auto* Y = Output(
|
0,
|
std::vector<int64_t>{N, out_h, out_w, kernel_h_ * kernel_w_ * C},
|
at::dtype<T>());
|
|
const size_t dx = X.numel() / N;
|
const size_t dy = Y->numel() / N;
|
for (int n = 0; n < N; ++n) {
|
const auto* xdata = X.template data<T>() + (n * dx);
|
auto* ydata = Y->template mutable_data<T>() + (n * dy);
|
math::Im2Col<T, Context, StorageOrder::NHWC>(
|
C,
|
H,
|
W,
|
kernel_h_,
|
kernel_w_,
|
dilation_h_,
|
dilation_w_,
|
pad_,
|
pad_,
|
pad_,
|
pad_,
|
stride_h_,
|
stride_w_,
|
xdata,
|
ydata,
|
&context_);
|
}
|
}; break;
|
default:
|
CAFFE_THROW("Unknown storage order: ", order_);
|
}
|
|
return true;
|
}
|
|
private:
|
int pad_;
|
int kernel_h_;
|
int kernel_w_;
|
int dilation_h_;
|
int dilation_w_;
|
int stride_h_;
|
int stride_w_;
|
StorageOrder order_;
|
};
|
|
template <typename T, class Context>
|
class Col2ImOp final : public Operator<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
template <class... Args>
|
explicit Col2ImOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
pad_(this->template GetSingleArgument<int>("pad", 0)),
|
kernel_h_(this->template GetSingleArgument<int>(
|
"kernel_h",
|
this->template GetSingleArgument<int>("kernel", 0))),
|
kernel_w_(this->template GetSingleArgument<int>(
|
"kernel_w",
|
this->template GetSingleArgument<int>("kernel", 0))),
|
dilation_h_(this->template GetSingleArgument<int>(
|
"dilation_h",
|
this->template GetSingleArgument<int>("dilation", 1))),
|
dilation_w_(this->template GetSingleArgument<int>(
|
"dilation_w",
|
this->template GetSingleArgument<int>("dilation", 1))),
|
stride_h_(this->template GetSingleArgument<int>(
|
"stride_h",
|
this->template GetSingleArgument<int>("stride", 1))),
|
stride_w_(this->template GetSingleArgument<int>(
|
"stride_w",
|
this->template GetSingleArgument<int>("stride", 1))),
|
order_(StringToStorageOrder(
|
this->template GetSingleArgument<string>("order", "NCHW"))) {
|
CAFFE_ENFORCE(kernel_h_ > 0);
|
CAFFE_ENFORCE(kernel_w_ > 0);
|
CAFFE_ENFORCE(dilation_h_ > 0);
|
CAFFE_ENFORCE(dilation_w_ > 0);
|
CAFFE_ENFORCE(stride_h_ > 0);
|
CAFFE_ENFORCE(stride_w_ > 0);
|
CAFFE_ENFORCE(pad_ >= 0);
|
}
|
|
bool RunOnDevice() override {
|
auto& X = Input(0);
|
auto& Z = Input(1);
|
|
auto* Y = Output(0, Z.sizes(), at::dtype<T>());
|
CAFFE_ENFORCE(4 == Y->dim());
|
|
int N = 0, C = 0, H = 0, W = 0;
|
switch (order_) {
|
case StorageOrder::NCHW:
|
N = Y->dim32(0);
|
C = Y->dim32(1);
|
H = Y->dim32(2);
|
W = Y->dim32(3);
|
break;
|
case StorageOrder::NHWC:
|
N = Y->dim32(0);
|
H = Y->dim32(1);
|
W = Y->dim32(2);
|
C = Y->dim32(3);
|
break;
|
default:
|
CAFFE_THROW("Unknown storage order: ", order_);
|
}
|
|
const int dkernel_h = dilation_h_ * (kernel_h_ - 1) + 1;
|
const int dkernel_w = dilation_w_ * (kernel_w_ - 1) + 1;
|
CAFFE_ENFORCE(H >= dkernel_h);
|
CAFFE_ENFORCE(W >= dkernel_w);
|
const int out_h = (H + 2 * pad_ - dkernel_h) / stride_h_ + 1;
|
const int out_w = (W + 2 * pad_ - dkernel_w) / stride_w_ + 1;
|
CAFFE_ENFORCE(X.numel() == N * kernel_h_ * kernel_w_ * C * out_h * out_w);
|
|
const size_t dx = X.numel() / N;
|
const size_t dy = Y->numel() / N;
|
|
// could template-specialize this, but it's test code...
|
switch (order_) {
|
case StorageOrder::NCHW: {
|
for (int n = 0; n < N; ++n) {
|
const auto* xdata = X.template data<T>() + (n * dx);
|
auto* ydata = Y->template mutable_data<T>() + (n * dy);
|
math::Col2Im<T, Context, StorageOrder::NCHW>(
|
C,
|
H,
|
W,
|
kernel_h_,
|
kernel_w_,
|
dilation_h_,
|
dilation_w_,
|
pad_,
|
pad_,
|
pad_,
|
pad_,
|
stride_h_,
|
stride_w_,
|
xdata,
|
ydata,
|
&context_);
|
}
|
}; break;
|
case StorageOrder::NHWC: {
|
for (int n = 0; n < N; ++n) {
|
const auto* xdata = X.template data<T>() + (n * dx);
|
auto* ydata = Y->template mutable_data<T>() + (n * dy);
|
math::Col2Im<T, Context, StorageOrder::NHWC>(
|
C,
|
H,
|
W,
|
kernel_h_,
|
kernel_w_,
|
dilation_h_,
|
dilation_w_,
|
pad_,
|
pad_,
|
pad_,
|
pad_,
|
stride_h_,
|
stride_w_,
|
xdata,
|
ydata,
|
&context_);
|
}
|
}; break;
|
default:
|
CAFFE_THROW("Unknown storage order: ", order_);
|
}
|
|
return true;
|
}
|
|
private:
|
int pad_;
|
int kernel_h_;
|
int kernel_w_;
|
int dilation_h_;
|
int dilation_w_;
|
int stride_h_;
|
int stride_w_;
|
StorageOrder order_;
|
};
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_OPERATORS_IM2COL_OP_H_
|
