Caffe2 - C++ API
A deep learning, cross platform ML framework
softmax_op.h
1 #ifndef CAFFE2_OPERATORS_SOFTMAX_OP_H_
2 #define CAFFE2_OPERATORS_SOFTMAX_OP_H_
3 
4 #include "caffe2/core/context.h"
5 #include "caffe2/core/logging.h"
6 #include "caffe2/core/operator.h"
7 #include "caffe2/utils/math.h"
8 
9 namespace caffe2 {
10 
11 template <typename T, class Context>
12 class SoftmaxOp final : public Operator<Context> {
13  public:
14  template <class... Args>
15  explicit SoftmaxOp(Args&&... args)
16  : Operator<Context>(std::forward<Args>(args)...),
17  axis_(this->template GetSingleArgument<int>("axis", 1)) {}
18  USE_OPERATOR_CONTEXT_FUNCTIONS;
19  bool RunOnDevice() override;
20 
21  protected:
22  int axis_;
23  Tensor scale_;
24  Tensor rowmax_;
25  Tensor sum_multiplier_;
26 };
27 
28 template <typename T, class Context>
29 class SoftmaxGradientOp final : public Operator<Context> {
30  public:
31  template <class... Args>
32  explicit SoftmaxGradientOp(Args&&... args)
33  : Operator<Context>(std::forward<Args>(args)...),
34  axis_(this->template GetSingleArgument<int>("axis", 1)) {}
35  USE_OPERATOR_CONTEXT_FUNCTIONS;
36  bool RunOnDevice() override;
37 
38  protected:
39  int axis_;
40  Tensor scale_;
41  Tensor sum_multiplier_;
42 };
43 
44 } // namespace caffe2
45 
46 #endif // CAFFE2_OPERATORS_SOFTMAX_OP_H_
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13