Caffe2 - C++ API
A deep learning, cross platform ML framework
im2col_op.cc
1 #include "caffe2/operators/im2col_op.h"
2 
3 namespace caffe2 {
4 REGISTER_CPU_OPERATOR(Im2Col, Im2ColOp<float, CPUContext>);
5 REGISTER_CPU_OPERATOR(Col2Im, Col2ImOp<float, CPUContext>);
6 
8  using GradientMakerBase::GradientMakerBase;
9  vector<OperatorDef> GetGradientDefs() override {
10  return SingleGradientDef(
11  "Col2Im",
12  "",
13  std::vector<string>{GO(0), I(0)},
14  std::vector<string>{GI(0)});
15  }
16 };
17 REGISTER_GRADIENT(Im2Col, GetIm2ColGradient);
18 
20  using GradientMakerBase::GradientMakerBase;
21  vector<OperatorDef> GetGradientDefs() override {
22  return SingleGradientDef(
23  "Im2Col", "", std::vector<string>{GO(0)}, std::vector<string>{GI(0)});
24  }
25 };
26 REGISTER_GRADIENT(Col2Im, GetCol2ImGradient);
27 
28 OPERATOR_SCHEMA(Im2Col)
29  .NumInputs(1)
30  .NumOutputs(1)
31  .SetDoc("The Im2Col operator from Matlab.")
32  .TensorInferenceFunction(
33  [](const OperatorDef& def, const vector<TensorShape>& in) {
34  ArgumentHelper helper(def);
35  auto pad = helper.GetSingleArgument<int>("pad", 0);
36  auto kernel_h = helper.GetSingleArgument<int>(
37  "kernel_h", helper.GetSingleArgument<int>("kernel", 0));
38  auto kernel_w = helper.GetSingleArgument<int>(
39  "kernel_w", helper.GetSingleArgument<int>("kernel", 0));
40  auto dilation_h = helper.GetSingleArgument<int>(
41  "dilation_h", helper.GetSingleArgument<int>("dilation", 1));
42  auto dilation_w = helper.GetSingleArgument<int>(
43  "dilation_w", helper.GetSingleArgument<int>("dilation", 1));
44  auto stride_h = helper.GetSingleArgument<int>(
45  "stride_h", helper.GetSingleArgument<int>("stride", 1));
46  auto stride_w = helper.GetSingleArgument<int>(
47  "stride_w", helper.GetSingleArgument<int>("stride", 1));
48  auto order = StringToStorageOrder(
49  helper.GetSingleArgument<string>("order", "NCHW"));
50 
51  const TensorShape& X = in[0];
52  int N = 0, C = 0, H = 0, W = 0;
53  switch (order) {
54  case StorageOrder::NCHW:
55  N = X.dims(0);
56  C = X.dims(1);
57  H = X.dims(2);
58  W = X.dims(3);
59  break;
60  case StorageOrder::NHWC:
61  N = X.dims(0);
62  H = X.dims(1);
63  W = X.dims(2);
64  C = X.dims(3);
65  break;
66  default:
67  CAFFE_THROW("Unknown storage order: ", order);
68  }
69 
70  const int dkernel_h = dilation_h * (kernel_h - 1) + 1;
71  const int dkernel_w = dilation_w * (kernel_w - 1) + 1;
72  CAFFE_ENFORCE(H >= dkernel_h);
73  CAFFE_ENFORCE(W >= dkernel_w);
74  const int out_h = (H + 2 * pad - dkernel_h) / stride_h + 1;
75  const int out_w = (W + 2 * pad - dkernel_w) / stride_w + 1;
76 
77  vector<TensorShape> out(1);
78  switch (order) {
79  case StorageOrder::NCHW:
80  out[0] = CreateTensorShape(
81  vector<int>{N, C * kernel_h * kernel_w, out_h, out_w},
82  TensorProto::FLOAT);
83  break;
84  case StorageOrder::NHWC:
85  out[0] = CreateTensorShape(
86  vector<int>{N, out_h, out_w, kernel_h * kernel_w * C},
87  TensorProto::FLOAT);
88  break;
89  default:
90  CAFFE_THROW("Unknown storage order: ", order);
91  }
92 
93  return out;
94  })
95  .Input(0, "X", "4-tensor in NCHW or NHWC.")
96  .Output(
97  0,
98  "Y",
99  "4-tensor. For NCHW: N x (C x kH x kW) x outH x outW."
100  "For NHWC: N x outH x outW x (kH x kW x C");
101 
102 OPERATOR_SCHEMA(Col2Im).NumInputs(2).NumOutputs(1);
103 
104 } // namespace caffe2
A helper class to index into arguments.
Definition: proto_utils.h:200
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
static vector< OperatorDef > SingleGradientDef(const Args &...args)
a helper function to allow one to create one single operator def, which is usually the case for many ...
Definition: static.cpp:64