Caffe2 - C++ API
A deep learning, cross platform ML framework
order_switch_ops.cc
1 
17 #include "caffe2/operators/order_switch_ops.h"
18 
19 namespace caffe2 {
20 
21 template <>
22 bool NHWC2NCHWOp<float, CPUContext>::RunOnDevice() {
23  auto& X = Input(0);
24  auto* Y = Output(0);
25  CAFFE_ENFORCE(X.ndim() == 4);
26  const int N = X.dim32(0), H = X.dim32(1), W = X.dim32(2), C = X.dim32(3);
27  Y->Resize(N, C, H, W);
28  const float* Xdata = X.data<float>();
29  float* Ydata = Y->mutable_data<float>();
30  for (int n = 0; n < N; ++n) {
31  for (int h = 0; h < H; ++h) {
32  for (int w = 0; w < W; ++w) {
33  for (int c = 0; c < C; ++c) {
34  Ydata[((n * C + c) * H + h) * W + w] = *(Xdata++);
35  }
36  }
37  }
38  }
39  return true;
40 }
41 
42 template <>
43 bool NCHW2NHWCOp<float, CPUContext>::RunOnDevice() {
44  auto& X = Input(0);
45  auto* Y = Output(0);
46  CAFFE_ENFORCE(X.ndim() == 4);
47  const int N = X.dim32(0), C = X.dim32(1), H = X.dim32(2), W = X.dim32(3);
48  Y->Resize(N, H, W, C);
49  const float* Xdata = X.data<float>();
50  float* Ydata = Y->mutable_data<float>();
51  for (int n = 0; n < N; ++n) {
52  for (int c = 0; c < C; ++c) {
53  for (int h = 0; h < H; ++h) {
54  for (int w = 0; w < W; ++w) {
55  Ydata[((n * H + h) * W + w) * C + c] = *(Xdata++);
56  }
57  }
58  }
59  }
60  return true;
61 }
62 
63 
64 REGISTER_CPU_OPERATOR(NHWC2NCHW, NHWC2NCHWOp<float, CPUContext>);
65 REGISTER_CPU_OPERATOR(NCHW2NHWC, NCHW2NHWCOp<float, CPUContext>);
66 
67 OPERATOR_SCHEMA(NHWC2NCHW)
68  .NumInputs(1)
69  .NumOutputs(1)
70  .TensorInferenceFunction([](const OperatorDef& /*unused*/ /*def*/,
71  const vector<TensorShape>& in) {
72  CAFFE_ENFORCE_EQ(
73  in[0].dims_size(), 4, "Input for NHWC2NCHW must be 4 dimensional");
74  vector<TensorShape> out(1);
75  out[0].add_dims(in[0].dims(0));
76  out[0].add_dims(in[0].dims(3));
77  out[0].add_dims(in[0].dims(1));
78  out[0].add_dims(in[0].dims(2));
79  return out;
80  })
81  .SetDoc(R"DOC(
82 The operator switches the order of data in a tensor from NHWC- sample index N,
83 height H, width H and channels C, to the NCHW order.
84 )DOC")
85  .Input(0, "data", "The input data (Tensor<float>) in the NHWC order.")
86  .Output(
87  0,
88  "output",
89  "The output tensor (Tensor<float>) in the NCHW order.");
90 
91 OPERATOR_SCHEMA(NCHW2NHWC).NumInputs(1).NumOutputs(1)
92  .SetDoc(R"DOC(
93 The operator switches the order of data in a tensor from NCHW- sample index N,
94 channels C, height H and width W, to the NHWC order.
95 )DOC")
96  .Input(0, "data", "The input data (Tensor<float>) in the NCHW order.")
97  .Output(0, "output", "The output tensor (Tensor<float>) in the NHWC order.");
98 
99 
101  using GradientMakerBase::GradientMakerBase;
102  vector<OperatorDef> GetGradientDefs() override {
103  return SingleGradientDef(
104  "NCHW2NHWC", "",
105  vector<string>{GO(0)},
106  vector<string>{GI(0)});
107  }
108 };
109 REGISTER_GRADIENT(NHWC2NCHW, GetNHWC2NCHWGradient);
110 
112  using GradientMakerBase::GradientMakerBase;
113  vector<OperatorDef> GetGradientDefs() override {
114  return SingleGradientDef(
115  "NHWC2NCHW", "",
116  vector<string>{GO(0)},
117  vector<string>{GI(0)});
118  }
119 };
120 REGISTER_GRADIENT(NCHW2NHWC, GetNCHW2NHWCGradient);
121 } // namespace caffe2
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 ...