Caffe2 - C++ API
A deep learning, cross platform ML framework
space_batch_op.cc
1 
17 #include "caffe2/operators/space_batch_op.h"
18 
19 namespace caffe2 {
20 
21 REGISTER_CPU_OPERATOR(SpaceToBatch, SpaceToBatchOp<CPUContext>);
22 OPERATOR_SCHEMA(SpaceToBatch).NumInputs(1).NumOutputs(1).SetDoc(R"DOC(
23 
24 SpaceToBatch for 4-D tensors of type T.
25 
26 Zero-pads and then rearranges (permutes) blocks of spatial data into
27 batch. More specifically, this op outputs a copy of the input tensor
28 where values from the height and width dimensions are moved to the
29 batch dimension. After the zero-padding, both height and width of the
30 input must be divisible by the block size.
31 
32 )DOC");
33 
34 REGISTER_CPU_OPERATOR(BatchToSpace, BatchToSpaceOp<CPUContext>);
35 OPERATOR_SCHEMA(BatchToSpace).NumInputs(1).NumOutputs(1).SetDoc(R"DOC(
36 
37 BatchToSpace for 4-D tensors of type T.
38 
39 Rearranges (permutes) data from batch into blocks of spatial data,
40 followed by cropping. This is the reverse transformation of
41 SpaceToBatch. More specifically, this op outputs a copy of the input
42 tensor where values from the batch dimension are moved in spatial
43 blocks to the height and width dimensions, followed by cropping along
44 the height and width dimensions.
45 
46 )DOC");
47 
49  using GradientMakerBase::GradientMakerBase;
50  vector<OperatorDef> GetGradientDefs() override {
51  return SingleGradientDef(
52  "BatchToSpace", "", vector<string>{GO(0)}, vector<string>{GI(0)});
53  }
54 };
55 
57  using GradientMakerBase::GradientMakerBase;
58  vector<OperatorDef> GetGradientDefs() override {
59  return SingleGradientDef(
60  "SpaceToBatch", "", vector<string>{GO(0)}, vector<string>{GI(0)});
61  }
62 };
63 REGISTER_GRADIENT(SpaceToBatch, GetSpaceToBatchGradient);
64 REGISTER_GRADIENT(BatchToSpace, GetBatchToSpaceGradient);
65 }
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 ...