1 #include "caffe2/operators/flatten_op.h" 5 REGISTER_CPU_OPERATOR(Flatten, FlattenOp<CPUContext>);
7 OPERATOR_SCHEMA(Flatten)
10 .TensorInferenceFunction([](
const OperatorDef& def,
11 const vector<TensorShape>& in) {
12 ArgumentHelper helper(def);
13 const int axis = helper.GetSingleArgument<
int>(
"axis", 1);
14 vector<TensorShape> out(1);
17 std::size_t index = 0;
18 for (
auto d : in[0].dims()) {
26 out[0].set_data_type(in[0].data_type());
27 out[0].add_dims(outer);
28 out[0].add_dims(inner);
32 Flattens the input tensor into a 2D matrix. If input tensor has shape 33 (d_0, d_1, ... d_n) then the output will have shape 34 (d_0 X d_1 ... d_(axis-1), d_axis X d_(axis+1) ... X dn) 36 .Input(0, "input",
"A tensor of rank >= axis.")
40 "A 2D tensor with the contents of the input tensor, " 41 "with input dimensions up to axis flattened to the outer dimension " 42 "of the output and remaining input dimensions flattened into the inner " 43 "dimension of the output.")
46 "(Default to 1) Indicate up to which input dimensions " 47 "(exclusive) should be flattened to the outer dimension of the output")
48 .InheritOnnxSchema(
"Flatten");
51 using GradientMakerBase::GradientMakerBase;
52 vector<OperatorDef> GetGradientDefs()
override {
54 "ResizeLike",
"", vector<string>{GO(0), I(0)}, vector<string>{GI(0)});
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
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 ...