Caffe2 - C++ API
A deep learning, cross platform ML framework
space_batch_op.h
1 
17 #ifndef CAFFE2_OPERATORS_SPACE_BATCH_OP_H_
18 #define CAFFE2_OPERATORS_SPACE_BATCH_OP_H_
19 
20 #include "caffe2/core/context.h"
21 #include "caffe2/core/logging.h"
22 #include "caffe2/core/operator.h"
23 #include "caffe2/utils/math.h"
24 
25 namespace caffe2 {
26 
27 template <typename Context>
28 void spaceToBatch(
29  const Tensor<Context>& input,
30  int pad_t,
31  int pad_l,
32  int block_size,
33  Tensor<Context>* output,
34  Context* /*context*/) {
35  CAFFE_ENFORCE(input.ndim() == 4);
36  CAFFE_ENFORCE(output->ndim() == 4);
37 
38  const int output_batch = output->dim32(0);
39  const int output_depth = output->dim32(1);
40  const int output_height = output->dim32(2);
41  const int output_width = output->dim32(3);
42 
43  const int input_batch = input.dim32(0);
44  const int input_depth = input.dim32(1);
45  const int input_height = input.dim32(2);
46  const int input_width = input.dim32(3);
47 
48  for (int out_b = 0; out_b < output_batch; ++out_b) {
49  const int in_b = out_b % input_batch;
50  const int offset_w = (out_b / input_batch) % block_size;
51  const int offset_h = (out_b / input_batch) / block_size;
52  for (int d = 0; d < input_depth; ++d) {
53  for (int out_h = 0; out_h < output_height; ++out_h) {
54  const int in_h = out_h * block_size + offset_h - pad_t;
55  for (int out_w = 0; out_w < output_width; ++out_w) {
56  const int in_w = out_w * block_size + offset_w - pad_l;
57  const auto output_offset =
58  ((out_b * output_depth + d) * output_height + out_h) *
59  output_width +
60  out_w;
61  const auto input_offset =
62  ((in_b * input_depth + d) * input_height + in_h) * input_width +
63  in_w;
64  if (in_h >= 0 && in_w >= 0 && in_h < input_height &&
65  in_w < input_width) {
66  output->template mutable_data<float>()[output_offset] =
67  input.template data<float>()[input_offset];
68  } else {
69  output->template mutable_data<float>()[output_offset] = 0.0;
70  }
71  }
72  }
73  }
74  }
75 }
76 
77 template <typename Context>
78 void batchToSpace(
79  const Tensor<Context>& input,
80  int pad_t,
81  int pad_l,
82  int block_size,
83  Tensor<Context>* output,
84  Context* /*context*/) {
85  CAFFE_ENFORCE(input.ndim() == 4);
86  CAFFE_ENFORCE(output->ndim() == 4);
87 
88  const int output_batch = output->dim32(0);
89  const int output_depth = output->dim32(1);
90  const int output_height = output->dim32(2);
91  const int output_width = output->dim32(3);
92 
93  const int input_batch = input.dim32(0);
94  const int input_depth = input.dim32(1);
95  const int input_height = input.dim32(2);
96  const int input_width = input.dim32(3);
97 
98  CAFFE_ENFORCE(input_depth == output_depth);
99  for (int in_b = 0; in_b < input_batch; ++in_b) {
100  const int out_b = in_b % output_batch;
101  const int offset_w = (in_b / output_batch) % block_size;
102  const int offset_h = (in_b / output_batch) / block_size;
103  for (int d = 0; d < input_depth; ++d) {
104  for (int in_h = 0; in_h < input_height; ++in_h) {
105  const int out_h = in_h * block_size + offset_h - pad_t;
106  for (int in_w = 0; in_w < input_width; ++in_w) {
107  const int out_w = in_w * block_size + offset_w - pad_l;
108  if (out_h >= 0 && out_w >= 0 && out_h < output_height &&
109  out_w < output_width) {
110  const auto output_offset =
111  ((out_b * output_depth + d) * output_height + out_h) *
112  output_width +
113  out_w;
114  const auto input_offset =
115  ((in_b * input_depth + d) * input_height + in_h) * input_width +
116  in_w;
117  output->template mutable_data<float>()[output_offset] =
118  input.template data<float>()[input_offset];
119  }
120  }
121  }
122  }
123  }
124 }
125 
126 template <typename Context>
127 class SpaceBatchOpBase : public Operator<Context> {
128  public:
129  USE_OPERATOR_CONTEXT_FUNCTIONS;
130  SpaceBatchOpBase(const OperatorDef& operator_def, Workspace* ws)
131  : Operator<Context>(operator_def, ws),
132  pad_(OperatorBase::GetSingleArgument<int>("pad", 0)),
133  pad_t_(OperatorBase::GetSingleArgument<int>("pad_t", pad_)),
134  pad_l_(OperatorBase::GetSingleArgument<int>("pad", pad_)),
135  pad_b_(OperatorBase::GetSingleArgument<int>("pad", pad_)),
136  pad_r_(OperatorBase::GetSingleArgument<int>("pad", pad_)),
137  block_size_(OperatorBase::GetSingleArgument<int>("block_size", 2)),
138  order_(StringToStorageOrder(
139  OperatorBase::GetSingleArgument<string>("order", "NCHW"))) {
140  CAFFE_ENFORCE(order_ == StorageOrder::NCHW);
141  }
142 
143  protected:
144  int pad_;
145  int pad_t_;
146  int pad_l_;
147  int pad_b_;
148  int pad_r_;
149  int block_size_;
150  StorageOrder order_;
151 };
152 
153 template <typename Context>
154 class SpaceToBatchOp final : public SpaceBatchOpBase<Context> {
155  public:
156  USE_OPERATOR_CONTEXT_FUNCTIONS;
158 
159  bool RunOnDevice() override {
160  const auto& input = Input(0);
161  auto* output = Output(0);
162  const int batch = input.dim32(0);
163  const int depth = input.dim32(1);
164  const int height = this->pad_b_ + this->pad_t_ + input.dim32(2);
165  const int width = this->pad_l_ + this->pad_r_ + input.dim32(3);
166  CAFFE_ENFORCE(
167  height % this->block_size_ == 0,
168  "Height: ",
169  height,
170  ", block size: ",
171  this->block_size_);
172  CAFFE_ENFORCE(width % this->block_size_ == 0);
173 
174  const int output_batch = batch * this->block_size_ * this->block_size_;
175  const int output_height = height / this->block_size_;
176  const int output_width = width / this->block_size_;
177  Output(0)->Resize(output_batch, depth, output_height, output_width);
178 
179  spaceToBatch<Context>(
180  input,
181  this->pad_t_,
182  this->pad_l_,
183  this->block_size_,
184  output,
185  &context_);
186 
187  return true;
188  }
189 };
190 
191 template <typename Context>
192 class BatchToSpaceOp final : public SpaceBatchOpBase<Context> {
193  public:
194  USE_OPERATOR_CONTEXT_FUNCTIONS;
196 
197  bool RunOnDevice() override {
198  const auto& input = Input(0);
199  auto* output = Output(0);
200  const int batch = input.dim32(0);
201  const int depth = input.dim32(1);
202  const int height = input.dim32(2);
203  const int width = input.dim32(3);
204 
205  const int output_batch = batch / this->block_size_ / this->block_size_;
206  const int output_height =
207  height * this->block_size_ - this->pad_b_ - this->pad_t_;
208  const int output_width =
209  width * this->block_size_ - this->pad_l_ - this->pad_r_;
210  Output(0)->Resize(output_batch, depth, output_height, output_width);
211  batchToSpace<Context>(
212  input,
213  this->pad_t_,
214  this->pad_l_,
215  this->block_size_,
216  output,
217  &context_);
218  return true;
219  }
220 };
221 
222 } // namespace caffe2
223 
224 #endif // CAFFE2_OPERATORS_SPACE_BATCH_OP_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.
Copyright (c) 2016-present, Facebook, Inc.