Caffe2 - C++ API
A deep learning, cross platform ML framework
prepend_dim_op.cc
1 
17 #include "caffe2/operators/prepend_dim_op.h"
18 
19 namespace caffe2 {
20 
21 REGISTER_CPU_OPERATOR(PrependDim, PrependDimOp<CPUContext>);
22 REGISTER_CPU_OPERATOR(MergeDim, MergeDimOp<CPUContext>);
23 
24 OPERATOR_SCHEMA(PrependDim)
25  .NumInputs(1)
26  .NumOutputs(1)
27  .AllowInplace({{0, 0}})
28  .SetDoc(R"DOC(
29 Reshape the tensor by prepending a dimension of fixed size and dividing the
30 size of the next dimension by that amount.
31 )DOC")
32  .Arg("dim_size", "Size of the dimension to prepend.")
33  .Input(0, "data", "An input tensor.")
34  .Output(0, "reshaped", "Reshaped tensor.");
35 
36 OPERATOR_SCHEMA(MergeDim)
37  .NumInputs(1)
38  .NumOutputs(1)
39  .AllowInplace({{0, 0}})
40  .SetDoc(R"DOC(
41 Merge first two dimensions in a single dimension with size dim(0) * dim(1).
42 )DOC")
43  .Input(0, "data", "An input tensor.")
44  .Output(0, "reshaped", "Reshaped tensor.");
45 
47  using GradientMakerBase::GradientMakerBase;
48  vector<OperatorDef> GetGradientDefs() override {
49  return SingleGradientDef(
50  "MergeDim", "", vector<string>{GO(0)}, vector<string>{GI(0)});
51  }
52 
53  // Arguments are no longer needed in backprop.
54  bool CopyArguments() const override {
55  return false;
56  }
57 };
58 
59 REGISTER_GRADIENT(PrependDim, GetPrependDimGradient);
60 
61 } // 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 ...