Caffe2 - C++ API
A deep learning, cross platform ML framework
im2col_op.cc
1 
17 #include "caffe2/operators/im2col_op.h"
18 
19 namespace caffe2 {
20 REGISTER_CPU_OPERATOR(Im2Col, Im2ColOp<float, CPUContext>);
21 REGISTER_CPU_OPERATOR(Col2Im, Col2ImOp<float, CPUContext>);
22 
24  using GradientMakerBase::GradientMakerBase;
25  vector<OperatorDef> GetGradientDefs() override {
26  return SingleGradientDef(
27  "Col2Im",
28  "",
29  std::vector<string>{O(0), I(0)},
30  std::vector<string>{GI(0)});
31  }
32 };
33 REGISTER_GRADIENT(Im2Col, GetIm2ColGradient);
34 
36  using GradientMakerBase::GradientMakerBase;
37  vector<OperatorDef> GetGradientDefs() override {
38  return SingleGradientDef(
39  "Im2Col", "", std::vector<string>{O(0)}, std::vector<string>{GI(0)});
40  }
41 };
42 REGISTER_GRADIENT(Col2Im, GetCol2ImGradient);
43 
44 OPERATOR_SCHEMA(Im2Col)
45  .NumInputs(1)
46  .NumOutputs(1)
47  .SetDoc("The Im2Col operator from Matlab.")
48  .TensorInferenceFunction(
49  [](const OperatorDef& def, const vector<TensorShape>& in) {
50  ArgumentHelper helper(def);
51  auto pad = helper.GetSingleArgument<int>("pad", 0);
52  auto kernel_h = helper.GetSingleArgument<int>(
53  "kernel_h", helper.GetSingleArgument<int>("kernel", 0));
54  auto kernel_w = helper.GetSingleArgument<int>(
55  "kernel_w", helper.GetSingleArgument<int>("kernel", 0));
56  auto dilation_h = helper.GetSingleArgument<int>(
57  "dilation_h", helper.GetSingleArgument<int>("dilation", 1));
58  auto dilation_w = helper.GetSingleArgument<int>(
59  "dilation_w", helper.GetSingleArgument<int>("dilation", 1));
60  auto stride_h = helper.GetSingleArgument<int>(
61  "stride_h", helper.GetSingleArgument<int>("stride", 1));
62  auto stride_w = helper.GetSingleArgument<int>(
63  "stride_w", helper.GetSingleArgument<int>("stride", 1));
64  auto order = StringToStorageOrder(
65  helper.GetSingleArgument<string>("order", "NCHW"));
66 
67  const TensorShape& X = in[0];
68  int N = 0, C = 0, H = 0, W = 0;
69  switch (order) {
70  case StorageOrder::NCHW:
71  N = X.dims(0);
72  C = X.dims(1);
73  H = X.dims(2);
74  W = X.dims(3);
75  break;
76  case StorageOrder::NHWC:
77  N = X.dims(0);
78  H = X.dims(1);
79  W = X.dims(2);
80  C = X.dims(3);
81  break;
82  default:
83  CAFFE_THROW("Unknown storage order: ", order);
84  }
85 
86  const int dkernel_h = dilation_h * (kernel_h - 1) + 1;
87  const int dkernel_w = dilation_w * (kernel_w - 1) + 1;
88  CAFFE_ENFORCE(H >= dkernel_h);
89  CAFFE_ENFORCE(W >= dkernel_w);
90  const int out_h = (H + 2 * pad - dkernel_h) / stride_h + 1;
91  const int out_w = (W + 2 * pad - dkernel_w) / stride_w + 1;
92 
93  vector<TensorShape> out(1);
94  switch (order) {
95  case StorageOrder::NCHW:
96  out[0] = CreateTensorShape(
97  vector<int>{N, C * kernel_h * kernel_w, out_h, out_w},
98  TensorProto::FLOAT);
99  break;
100  case StorageOrder::NHWC:
101  out[0] = CreateTensorShape(
102  vector<int>{N, out_h, out_w, kernel_h * kernel_w * C},
103  TensorProto::FLOAT);
104  break;
105  default:
106  CAFFE_THROW("Unknown storage order: ", order);
107  }
108 
109  return out;
110  })
111  .Input(0, "X", "4-tensor in NCHW or NHWC.")
112  .Output(
113  0,
114  "Y",
115  "4-tensor. For NCHW: N x (C x kH x kW) x outH x outW."
116  "For NHWC: N x outH x outW x (kH x kW x C");
117 
118 OPERATOR_SCHEMA(Col2Im).NumInputs(2).NumOutputs(1);
119 
120 } // namespace caffe2
A helper class to index into arguments.
Definition: proto_utils.h:198
Copyright (c) 2016-present, Facebook, Inc.
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 ...