Caffe2 - C++ API
A deep learning, cross platform ML framework
softmax_op.h
1 
17 #ifndef CAFFE2_OPERATORS_SOFTMAX_OP_H_
18 #define CAFFE2_OPERATORS_SOFTMAX_OP_H_
19 
20 #include "caffe2/core/context.h"
21 #include "caffe2/core/logging.h"
22 #include "caffe2/core/operator.h"
23 #include "caffe2/utils/math.h"
24 
25 namespace caffe2 {
26 
27 template <typename T, class Context>
28 class SoftmaxOp final : public Operator<Context> {
29  public:
30  SoftmaxOp(const OperatorDef& operator_def, Workspace* ws)
31  : Operator<Context>(operator_def, ws),
32  axis_(OperatorBase::GetSingleArgument<int>("axis", 1)) {}
33  USE_OPERATOR_CONTEXT_FUNCTIONS;
34  bool RunOnDevice() override;
35 
36  protected:
37  int axis_;
38  Tensor<Context> scale_;
39  Tensor<Context> rowmax_;
40  Tensor<Context> sum_multiplier_;
41 };
42 
43 template <typename T, class Context>
44 class SoftmaxGradientOp final : public Operator<Context> {
45  public:
46  SoftmaxGradientOp(const OperatorDef& operator_def, Workspace* ws)
47  : Operator<Context>(operator_def, ws),
48  axis_(OperatorBase::GetSingleArgument<int>("axis", 1)) {}
49  USE_OPERATOR_CONTEXT_FUNCTIONS;
50  bool RunOnDevice() override;
51 
52  protected:
53  int axis_;
54  Tensor<Context> scale_;
55  Tensor<Context> sum_multiplier_;
56 };
57 
58 } // namespace caffe2
59 
60 #endif // CAFFE2_OPERATORS_SOFTMAX_OP_H_
Tensor is the basic class in Caffe2 that stores a contiguous memory with its shape information...
Definition: tensor.h:109
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.