forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathspace_batch_op.cc
123 lines (88 loc) · 3.92 KB
/
space_batch_op.cc
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);
}