1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
|
#include "caffe2/operators/space_batch_op.h"
namespace caffe2 {
REGISTER_CPU_OPERATOR(SpaceToBatch, SpaceToBatchOp<CPUContext>);
OPERATOR_SCHEMA(SpaceToBatch).NumInputs(1).NumOutputs(1).SetDoc(R"DOC(
Zero-pads and then rearranges (permutes) blocks of spatial data into batch. More specifically, this op outputs a copy of the input tensor where values from the height and width dimensions are moved to the batch dimension. After the zero-padding is according to the `pad` argument, both height and width of the input must be divisible by the `block_size`. Only "NCHW" order is currently supported.
Github Links:
- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/space_batch_op.cc
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"SpaceToBatch",
["X"],
["Y"],
pad=2,
block_size=3
)
workspace.FeedBlob("X", np.random.rand(1,3,5,5).astype(np.float32))
print("X.shape:", workspace.FetchBlob("X").shape)
workspace.RunOperatorOnce(op)
print("Y.shape:", workspace.FetchBlob("Y").shape)
```
**Result**
```
X.shape: (1, 3, 5, 5)
Y.shape: (9, 3, 3, 3)
```
</details>
)DOC")
.Arg("pad","(*int*): exclusive axis that divides the first and second dimension of matrix `A` (default=0)")
.Arg("block_size","(*int*): height/width of spatial blocks to be moved (default=2)")
// NOLINTNEXTLINE(modernize-raw-string-literal)
.Arg("order","(*string*): order of dimensions of input and output blobs; only \"NCHW\" order is currently supported (default=\"NCHW\")")
.Input(0,"X","(*Tensor`<float>`*): input tensor (NCHW order)")
.Output(0,"Y","(*Tensor`<float>`*): output tensor (NCHW order)");
REGISTER_CPU_OPERATOR(BatchToSpace, BatchToSpaceOp<CPUContext>);
OPERATOR_SCHEMA(BatchToSpace).NumInputs(1).NumOutputs(1).SetDoc(R"DOC(
Rearranges (permutes) data from batch into blocks of spatial data, followed by cropping. This is the reverse transformation of `SpaceToBatch`. More specifically, this op outputs a copy of the input tensor where values from the batch dimension are moved in spatial blocks to the height and width dimensions, followed by cropping along the height and width dimensions. Only "NCHW" order is currently supported.
Github Links:
- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/space_batch_op.cc
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"BatchToSpace",
["X"],
["Y"],
pad=3
)
workspace.FeedBlob("X", np.random.rand(10,3,32,32).astype(np.float32))
print("X.shape:", workspace.FetchBlob("X").shape)
workspace.RunOperatorOnce(op)
print("Y.shape:", workspace.FetchBlob("Y").shape)
```
**Result**
```
X.shape: (10, 3, 32, 32)
Y.shape: (2, 3, 58, 58)
```
</details>
)DOC")
.Arg("pad","(*int*): exclusive axis that divides the first and second dimension of matrix `A` (default=0)")
.Arg("block_size","(*int*): height/width of spatial blocks to be moved (default=2)")
// NOLINTNEXTLINE(modernize-raw-string-literal)
.Arg("order","(*string*): order of dimensions of input and output blobs; only \"NCHW\" order is currently supported (default=\"NCHW\")")
.Input(0,"X","(*Tensor`<float>`*): input tensor (NCHW order)")
.Output(0,"Y","(*Tensor`<float>`*): output tensor (NCHW order)");
class GetSpaceToBatchGradient : public GradientMakerBase {
using GradientMakerBase::GradientMakerBase;
vector<OperatorDef> GetGradientDefs() override {
return SingleGradientDef(
"BatchToSpace", "", vector<string>{GO(0)}, vector<string>{GI(0)});
}
};
class GetBatchToSpaceGradient : public GradientMakerBase {
using GradientMakerBase::GradientMakerBase;
vector<OperatorDef> GetGradientDefs() override {
return SingleGradientDef(
"SpaceToBatch", "", vector<string>{GO(0)}, vector<string>{GI(0)});
}
};
REGISTER_GRADIENT(SpaceToBatch, GetSpaceToBatchGradient);
REGISTER_GRADIENT(BatchToSpace, GetBatchToSpaceGradient);
}
|