Caffe2 - C++ API
A deep learning, cross platform ML framework
lengths_top_k_op.h
1 
18 #ifndef CAFFE2_OPERATORS_LENGTHS_TOP_K_OP_H_
19 #define CAFFE2_OPERATORS_LENGTHS_TOP_K_OP_H_
20 
21 #include "caffe2/core/context.h"
22 #include "caffe2/core/logging.h"
23 #include "caffe2/core/operator.h"
24 #include "caffe2/operators/conv_pool_op_base.h"
25 #include "caffe2/utils/math.h"
26 
27 namespace caffe2 {
28 template <typename T, class Context>
29 class LengthsTopKOp : public Operator<Context> {
30  public:
31  USE_OPERATOR_CONTEXT_FUNCTIONS;
32 
33  LengthsTopKOp(const OperatorDef& operator_def, Workspace* ws)
34  : Operator<Context>(operator_def, ws), OP_SINGLE_ARG(int, "k", k_, -1) {
35  CAFFE_ENFORCE_GE(k_, 1, "k argument must be >= 1");
36  }
37 
38  bool RunOnDevice() override;
39 
40  protected:
41  int k_;
42  INPUT_TAGS(X_IN, Y_IN);
43  OUTPUT_TAGS(TOPK_VALUES_OUT, TOPK_INDICES_OUT);
44 };
45 
46 template <typename T, class Context>
47 class LengthsTopKGradientOp : public Operator<Context> {
48  public:
49  LengthsTopKGradientOp(const OperatorDef& def, Workspace* ws)
50  : Operator<Context>(def, ws), OP_SINGLE_ARG(int, "k", k_, -1) {
51  CAFFE_ENFORCE_GE(k_, 1, "k argument must be >= 1");
52  }
53  USE_OPERATOR_CONTEXT_FUNCTIONS;
54 
55  bool RunOnDevice() override;
56 
57  protected:
58  int k_;
59  INPUT_TAGS(LENGTH_IN, INDICES_IN, DER_TOPK_IN);
60  OUTPUT_TAGS(DER_X_OUT);
61 };
62 
63 } // namespace caffe2
64 
65 #endif // CAFFE2_OPERATORS_LENGTHS_TOP_K_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.