Caffe2 - C++ API
A deep learning, cross platform ML framework
elementwise_linear_op.h
1 
17 #ifndef CAFFE2_OPERATORS_ELEMENTWISE_LINEAR_OP_H_
18 #define CAFFE2_OPERATORS_ELEMENTWISE_LINEAR_OP_H_
19 
20 #include "caffe2/core/context.h"
21 #include "caffe2/core/operator.h"
22 #include "caffe2/utils/math.h"
23 
24 namespace caffe2 {
25 template <typename T, class Context, class Engine = DefaultEngine>
26 class ElementwiseLinearOp final : public Operator<Context> {
27  public:
28  ElementwiseLinearOp(const OperatorDef& operator_def, Workspace* ws)
29  : Operator<Context>(operator_def, ws),
30  axis_(OperatorBase::GetSingleArgument<int>("axis", 1)) {}
31 
32  USE_OPERATOR_CONTEXT_FUNCTIONS;
33  bool RunOnDevice() override;
34 
35  protected:
36  int axis_;
37 };
38 
39 template <typename T, class Context, class Engine = DefaultEngine>
40 class ElementwiseLinearGradientOp final : public Operator<Context> {
41  public:
42  ElementwiseLinearGradientOp(const OperatorDef& operator_def, Workspace* ws)
43  : Operator<Context>(operator_def, ws),
44  axis_(OperatorBase::GetSingleArgument<int>("axis", 1)) {}
45 
46  USE_OPERATOR_CONTEXT_FUNCTIONS;
47  bool RunOnDevice() override;
48 
49  protected:
50  int axis_;
51 };
52 
53 } // namespace caffe2
54 
55 #endif // CAFFE2_OPERATORS_ELEMENTWISE_LINEAR_OP_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.