Caffe2 - C++ API
A deep learning, cross platform ML framework
flexible_top_k.h
1 // Copyright 2004-present Facebook. All Rights Reserved.
2 
3 #ifndef CAFFE2_OPERATORS_FLEXIBLE_TOP_K_H_
4 #define CAFFE2_OPERATORS_FLEXIBLE_TOP_K_H_
5 
6 #include "caffe2/core/logging.h"
7 #include "caffe2/core/operator.h"
8 #include "caffe2/utils/math.h"
9 
10 namespace caffe2 {
11 
12 template <typename T, class Context>
13 class FlexibleTopKOp : public Operator<Context> {
14  public:
15  USE_OPERATOR_CONTEXT_FUNCTIONS;
16 
17  FlexibleTopKOp(const OperatorDef& operator_def, Workspace* ws)
18  : Operator<Context>(operator_def, ws) {}
19 
20  bool RunOnDevice() override;
21 };
22 
23 template <typename T, class Context>
24 class FlexibleTopKGradientOp : public Operator<Context> {
25  public:
26  USE_OPERATOR_CONTEXT_FUNCTIONS;
27 
28  FlexibleTopKGradientOp(const OperatorDef& operator_def, Workspace* ws)
29  : Operator<Context>(operator_def, ws) {}
30 
31  bool RunOnDevice() override;
32 };
33 
34 } // namespace caffe2
35 
36 #endif // CAFFE2_OPERATORS_FLEXIBLE_TOP_K_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.