Caffe2 - C++ API
A deep learning, cross platform ML framework
fully_connected_op.cc
1 
17 #include <functional>
18 
19 #include "caffe2/operators/fully_connected_op.h"
20 
21 namespace caffe2 {
22 
23 REGISTER_CPU_OPERATOR(FC, FullyConnectedOp<CPUContext>);
24 REGISTER_CPU_OPERATOR(FCGradient, FullyConnectedGradientOp<CPUContext>);
25 
26 REGISTER_CPU_OPERATOR(
27  FCTransposed,
28  FullyConnectedOp<
29  CPUContext,
30  DefaultEngine,
31  false /* don't transpose weight */>);
32 REGISTER_CPU_OPERATOR(
33  FCTransposedGradient,
34  FullyConnectedGradientOp<
35  CPUContext,
36  DefaultEngine,
37  false /* don't transpose weight */>);
38 
39 namespace {
40 std::vector<TensorShape> FCShapeInference(
41  const OperatorDef& def,
42  const vector<TensorShape>& in,
43  bool pretransposed_weight) {
44  vector<TensorShape> out(1);
45  ArgumentHelper helper(def);
46 
47  auto axis = helper.GetSingleArgument<int32_t>("axis", 1);
48  const auto canonical_axis = canonical_axis_index_(axis, in[0].dims().size());
49  const int M = size_to_dim_(canonical_axis, GetDimsVector(in[0]));
50  auto axis_w = helper.GetSingleArgument<int32_t>("axis_w", 1);
51  const int canonical_axis_w =
52  canonical_axis_index_(axis_w, in[1].dims().size());
53  const int N = pretransposed_weight
54  ? size_from_dim_(canonical_axis_w, GetDimsVector(in[1]))
55  : size_to_dim_(canonical_axis_w, GetDimsVector(in[1]));
56 
57  vector<int> y_shape(in[0].dims().begin(), in[0].dims().end());
58  CAFFE_ENFORCE_LE(canonical_axis + 1, y_shape.size());
59  y_shape.resize(canonical_axis + 1);
60  y_shape[canonical_axis] = N;
61  out[0] = CreateTensorShape(y_shape, in[0].data_type());
62  return out;
63 }
64 
65 OpSchema::Cost CostInferenceForFC(
66  const OperatorDef& def,
67  const vector<TensorShape>& in) {
68  struct OpSchema::Cost c;
69  ArgumentHelper helper(def);
70 
71  auto axis = helper.GetSingleArgument<int32_t>("axis", 1);
72  const auto canonical_axis = canonical_axis_index_(axis, in[0].dims().size());
73  const int M = size_to_dim_(canonical_axis, GetDimsVector(in[0]));
74  const int K = size_from_dim_(canonical_axis, GetDimsVector(in[0]));
75  auto axis_w = helper.GetSingleArgument<int32_t>("axis_w", 1);
76  const int canonical_axis_w =
77  canonical_axis_index_(axis_w, in[1].dims().size());
78  const int N = size_to_dim_(canonical_axis_w, GetDimsVector(in[1]));
79  c.flops = 2 * K * M * N + M * N;
80  c.bytes_moved = M * N * sizeof(float);
81  c.params_bytes = (K * N + N) * sizeof(float);
82  return c;
83 }
84 } // namespace
85 
86 using namespace std::placeholders;
87 OPERATOR_SCHEMA(FCTransposed)
88  .NumInputs(3)
89  .NumOutputs(1)
90  .TensorInferenceFunction(std::bind(FCShapeInference, _1, _2, true))
91  .SetDoc(R"DOC(
92 Same as FC, but weight matrix is supposed to be already pretransposed.
93 FCTransposed stands for calling blass with no noTrans, noTrans
94 )DOC");
95 
96 OPERATOR_SCHEMA(FC)
97  .NumInputs(3)
98  .NumOutputs(1)
99  .TensorInferenceFunction(std::bind(FCShapeInference, _1, _2, false))
100  .CostInferenceFunction(
101  OpSchema::CostInferenceFunctionType(CostInferenceForFC))
102  .SetDoc(R"DOC(
103 Computes the result of passing an input vector X into a fully
104 connected layer with 2D weight matrix W and 1D bias vector b. That is,
105 the layer computes Y = X * W^T + b, where X has size (M x K),
106 W has size (N x K), b has size (N), and Y has size (M x N),
107 where M is often the batch size.
108 
109 
110 NOTE: X does not need to explicitly be a 2D vector; rather, it will be
111 coerced into one. For an arbitrary n-dimensional tensor
112 X \in [a_0, a_1, ...,a_{k-1}, a_k, ..., a_{n-1}] where a_i \in N+ and k is
113 the axis provided, then X will be coerced into a 2-dimensional tensor with
114 dimensions [a_0 * ... * a_{k-1}, a_k * ... * a_{n-1}]. For the default
115 case where axis=1, this means the X tensor will be coerced into a 2D tensor
116 of dimensions [a_0, a_1 * ... * a_{n-1}], where a_0 is often the batch size.
117 In this situation, we must have a_0 = M and a_1 * ... * a_{n-1} = K.
118 Lastly, even though b is a 1D vector of size N, it is copied/resized to
119 be size (M x N) implicitly and added to each vector in the batch.
120 Each of these dimensions must be matched correctly, or else the operator
121 will throw errors.
122 )DOC")
123  .Arg(
124  "axis",
125  "(int32_t) default to 1; describes the axis of the inputs; "
126  "defaults to one because the 0th axis most likely describes "
127  "the batch_size")
128  .Arg(
129  "axis_w",
130  "(int32_t) default to 1; describes the axis of the weight matrix W; "
131  "defaults to one because the 0th axis most likely describes "
132  "the batch_size")
133  .Arg("float16_compute", "Whether to use float-16 compute kernel")
134  .Input(
135  0,
136  "X",
137  "input tensor that's coerced into a 2D matrix of size (MxK) "
138  "as described above")
139  .Input(
140  1,
141  "W",
142  "A tensor that is coerced into a 2D blob of size (KxN) "
143  "containing fully connected weight matrix")
144  .Input(2, "b", "1D blob containing bias vector")
145  .Output(0, "Y", "2D output tensor");
146 
147 OPERATOR_SCHEMA(FCGradient).NumInputs(3).NumOutputs(2, 3);
148 OPERATOR_SCHEMA(FCTransposedGradient).NumInputs(3).NumOutputs(2, 3);
149 
150 namespace {
151 
152 class GetFCGradient : public GradientMakerBase {
153  using GradientMakerBase::GradientMakerBase;
154 
155  std::vector<OperatorDef> GetGradientDefs() override {
156  CAFFE_ENFORCE_EQ(def_.input_size(), 3);
157  CAFFE_ENFORCE(def_.type() == "FC" || def_.type() == "FCTransposed");
158  return SingleGradientDef(
159  def_.type() + "Gradient",
160  "",
161  vector<string>{I(0), I(1), GO(0)},
162  vector<string>{GI(1), GI(2), GI(0)});
163  }
164 };
165 
166 REGISTER_GRADIENT(FC, GetFCGradient);
167 REGISTER_GRADIENT(FCTransposed, GetFCGradient);
168 
169 } // namespace
170 
171 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.
std::function< struct Cost(const OperatorDef &, const vector< TensorShape > &)> CostInferenceFunctionType
Registers a function that takes in an OperatorDef and a series of input shapes and returns the total ...
TIndex size_from_dim_(int k, vector< TIndex > dims)
Return product of all dimensions starting from K.
Definition: tensor.h:56