Caffe2 - C++ API
A deep learning, cross platform ML framework
flatten_op.cc
1 
17 #include "caffe2/operators/flatten_op.h"
18 
19 namespace caffe2 {
20 
21 REGISTER_CPU_OPERATOR(Flatten, FlattenOp<CPUContext>);
22 
23 OPERATOR_SCHEMA(Flatten)
24  .NumInputs(1)
25  .NumOutputs(1)
26  .TensorInferenceFunction([](const OperatorDef& def,
27  const vector<TensorShape>& in) {
28  ArgumentHelper helper(def);
29  const int axis = helper.GetSingleArgument<int>("axis", 1);
30  vector<TensorShape> out(1);
31  TIndex outer = 1;
32  TIndex inner = 1;
33  std::size_t index = 0;
34  for (auto d : in[0].dims()) {
35  if (index < axis) {
36  outer *= d;
37  } else {
38  inner *= d;
39  }
40  ++index;
41  }
42  out[0].set_data_type(in[0].data_type());
43  out[0].add_dims(outer);
44  out[0].add_dims(inner);
45  return out;
46  })
47  .SetDoc(R"DOC(
48 Flattens the input tensor into a 2D matrix. If input tensor has shape
49 (d_0, d_1, ... d_n) then the output will have shape
50 (d_0 X d_1 ... d_(axis-1), d_axis X d_(axis+1) ... X dn)
51 )DOC")
52  .Input(0, "input", "A tensor of rank >= axis.")
53  .Output(
54  0,
55  "output",
56  "A 2D tensor with the contents of the input tensor, "
57  "with input dimensions up to axis flattened to the outer dimension "
58  "of the output and remaining input dimensions flattened into the inner "
59  "dimension of the output.")
60  .Arg(
61  "axis",
62  "(Default to 1) Indicate up to which input dimensions "
63  "(exclusive) should be flattened to the outer dimension of the output");
64 
66  using GradientMakerBase::GradientMakerBase;
67  vector<OperatorDef> GetGradientDefs() override {
68  return SingleGradientDef(
69  "ResizeLike", "", vector<string>{GO(0), I(0)}, vector<string>{GI(0)});
70  }
71 };
72 
73 REGISTER_GRADIENT(Flatten, GetFlattenGradient);
74 
75 } // namespace caffe2
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 ...