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
124
125
126
127
128
129
130
131
132
133
134
|
#include <caffe2/ideep/ideep_utils.h>
using namespace caffe2;
namespace {
// Takes a shape and data tensor and reshapes it
class IDEEPReshapeOp final : public IDEEPOperator {
public:
USE_IDEEP_DEF_ALIASES();
USE_IDEEP_OPERATOR_FUNCTIONS();
IDEEPReshapeOp(const OperatorDef& operator_def, Workspace* ws)
: IDEEPOperator(operator_def, ws),
new_shape_(OperatorBase::GetRepeatedArgument<itensor::dim>("shape")) {}
bool RunOnDevice() override {
ideep::tensor::dims actual_new_shape = new_shape_;
if (InputSize() == 2) {
CAFFE_ENFORCE(
!OperatorBase::HasArgument("shape"),
"New shape is specified by the input blob, do not pass in "
"the argument `shape`.");
// shape info live on CPU
auto& shape = OperatorBase::Input<TensorCPU>(1, CPU);
CAFFE_ENFORCE(shape.ndim() == 1, "Shape should be 1-D");
actual_new_shape.reserve(shape.size());
if (shape.template IsType<int>()) {
const int* shape_data = shape.template data<int>();
actual_new_shape.assign(shape_data, shape_data + shape.size());
} else if (shape.template IsType<int64_t>()) {
const int64_t* shape_data = shape.template data<int64_t>();
for (int i = 0; i < shape.size(); ++i) {
actual_new_shape.push_back(static_cast<int>(shape_data[i]));
}
} else {
CAFFE_THROW(
"IDEEP reshape only supports shape data in int32_t or int64_t");
}
} else {
CAFFE_ENFORCE(
OperatorBase::HasArgument("shape"), "Argument `shape` is missing.");
}
auto& input = Input(0);
// Copy over the dimensions for those that are specified zero.
// NOLINTNEXTLINE(clang-diagnostic-sign-compare)
for (int i = 0; i < actual_new_shape.size() && i < input.ndims(); ++i) {
if (actual_new_shape[i] == 0) {
actual_new_shape[i] = input.get_dim(i);
}
}
// Checks if the new shape is valid and fills in the missing dimension
// specified by -1.
// NOTE: At most one dimension can be -1.
auto total_size = input.get_nelems();
int size = 1;
int unknown_idx = -1;
// NOLINTNEXTLINE(clang-diagnostic-sign-compare)
for (int i = 0; i < actual_new_shape.size(); ++i) {
const auto dim = actual_new_shape[i];
if (dim == -1) {
CAFFE_ENFORCE(
unknown_idx == -1,
"Argument `shape` has more than one missing dimension.");
unknown_idx = i;
} else {
size *= dim;
}
}
if (size == 0 && total_size != 0) {
CAFFE_THROW(
"Can not reshape a non-zero size (",
total_size,
") tensor to zero size.");
}
if (unknown_idx != -1) {
CAFFE_ENFORCE_NE(
size,
0,
"New shape at dim ",
unknown_idx,
" can not be inferred since new size is zero.");
CAFFE_ENFORCE(
total_size % size == 0,
"Argument `shape` does not agree with the input data.",
" (",
total_size,
" vs ",
size,
")");
actual_new_shape[unknown_idx] = total_size / size;
} else {
CAFFE_ENFORCE_EQ(
total_size,
size,
"Argument `shape` does not agree with the input data.",
" (",
total_size,
" != ",
size,
")");
}
// Write the original shape to the second output.
// shape info live on CPU
TensorCPU* old_shape = OperatorBase::Output<TensorCPU>(1, CPU);
old_shape->Resize(input.ndims());
int* old_shape_data = old_shape->template mutable_data<int>();
for (int i = 0; i < input.ndims(); ++i) {
old_shape_data[i] = input.get_dim(i);
}
auto* output = Output(0);
if (output != &input) {
// If we are not doing in-place computation, a copy is needed.
output->reinit_like(input);
ideep::direct_copy::compute(input, *output);
}
output->reshape(actual_new_shape);
return true;
}
private:
ideep::tensor::dims new_shape_;
};
REGISTER_IDEEP_OPERATOR(Reshape, IDEEPReshapeOp);
} // namespace
|
