Caffe2 - C++ API
A deep learning, cross platform ML framework
conv_op.h
1 #ifndef CAFFE2_OPERATORS_CONV_OP_H_
2 #define CAFFE2_OPERATORS_CONV_OP_H_
3 
4 #include "caffe2/core/context.h"
5 #include "caffe2/core/operator.h"
6 #include "caffe2/operators/conv_op_shared.h"
7 #include "caffe2/operators/conv_pool_op_base.h"
8 
9 C10_DECLARE_bool(caffe2_force_shared_col_buffer);
10 
11 namespace caffe2 {
12 
13 template <typename T, class Context>
14 class ConvOp final : public ConvPoolOpBase<Context> {
15  public:
16  USE_CONV_POOL_BASE_FUNCTIONS(Context);
17  explicit ConvOp(const OperatorDef& operator_def, Workspace* ws)
18  : ConvPoolOpBase<Context>(operator_def, ws) {
19  // Since this is the default convolution implementation, we will
20  // use CAFFE_ENFORCE instead of OPERATOR_NEEDS_FEATURE.
21  CAFFE_ENFORCE(
22  (group_ == 1 || order_ == StorageOrder::NCHW ||
23  std::is_same<Context, CPUContext>::value),
24  "Group convolution only supports NCHW order or CPUContext right now.");
25 
26  // Create shared buffer mutex in the constructor
27  // to avoid race-condition in DAGNet.
28  if (FLAGS_caffe2_force_shared_col_buffer || shared_buffer_) {
29  createSharedBuffer<Context>(ws_);
30  }
31  }
32  ~ConvOp() {}
33 
34  bool RunOnDeviceWithOrderNCHW() override;
35  bool RunOnDeviceWithOrderNHWC() override;
36 
37  private:
38  bool Run1x1ConvOnDeviceWithOrderNCHW(
39  const int N,
40  const int C,
41  const int HxW,
42  const int M,
43  const T* X,
44  const T* filter,
45  const T* bias,
46  T* Y);
47 
48  bool Run1x1ConvOnDeviceWithOrderNHWC(
49  const int N,
50  const int C,
51  const int HxW,
52  const int M,
53  const T* X,
54  const T* filter,
55  const T* bias,
56  T* Y);
57 
58  Tensor col_buffer_{Context::GetDeviceType()};
59  Tensor bias_multiplier_{Context::GetDeviceType()};
60  Tensor img_shape_device_{Context::GetDeviceType()};
61  Tensor col_buffer_shape_device_{Context::GetDeviceType()};
62  // Input: X, W, b
63  // Output: Y
64  INPUT_TAGS(INPUT, FILTER, BIAS);
65 };
66 
67 template <typename T, class Context>
68 class ConvGradientOp final : public ConvPoolOpBase<Context> {
69  public:
70  USE_CONV_POOL_BASE_FUNCTIONS(Context);
71  explicit ConvGradientOp(const OperatorDef& operator_def, Workspace* ws)
72  : ConvPoolOpBase<Context>(operator_def, ws),
73  no_bias_(this->template GetSingleArgument<int>("no_bias", 0)) {
74  CAFFE_ENFORCE(
75  !(no_bias_ && OutputSize() == 3),
76  "If bias is not present, you should not have 3 grad output.");
77  CAFFE_ENFORCE(
78  (group_ == 1 || order_ == StorageOrder::NCHW ||
79  std::is_same<Context, CPUContext>::value),
80  "Group convolution only supports NCHW order or CPUContext right now.");
81  }
82  ~ConvGradientOp() {}
83 
84  bool RunOnDeviceWithOrderNCHW() override;
85  bool RunOnDeviceWithOrderNHWC() override;
86 
87  private:
88  Tensor col_buffer_;
89  Tensor bias_multiplier_;
90  Tensor img_shape_device_{Context::GetDeviceType()};
91  Tensor col_buffer_shape_device_{Context::GetDeviceType()};
92  bool no_bias_;
93  // input: X, W, dY
94  // output: dW, db, and optionally dX
95  INPUT_TAGS(INPUT, FILTER, OUTPUT_GRAD);
96  OUTPUT_TAGS(FILTER_GRAD, BIAS_OR_INPUT_GRAD, INPUT_GRAD);
97 };
98 
99 } // namespace caffe2
100 
101 #endif // CAFFE2_OPERATORS_CONV_OP_H_
Definition: any.cpp:108
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:47
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
Definition: static.cpp:64