Dragon - C++ API
A Computation Graph Virtual Machine Based Deep Learning Framework
roi_pool_op.h
Go to the documentation of this file.
1 
13 #ifndef DRAGON_OPERATORS_VISION_ROI_POOL_OP_H_
14 #define DRAGON_OPERATORS_VISION_ROI_POOL_OP_H_
15 
16 #include "core/operator.h"
17 
18 namespace dragon {
19 
20 template <class Context>
21 class ROIPoolOp final : public Operator<Context> {
22  public:
23  ROIPoolOp(const OperatorDef& def, Workspace *ws)
24  : Operator<Context>(def, ws),
25  pool_h_(OpArg<int64_t>("pool_h", 0)),
26  pool_w_(OpArg<int64_t>("pool_w", 0)),
27  spatial_scale_(OpArg<float>("spatial_scale", 1.f)) {}
29 
30  void RunOnDevice() override;
31  template <typename T> void RunImpl();
32 
33  protected:
35  int64_t pool_h_, pool_w_;
36 };
37 
38 template <class Context>
39 class ROIPoolGradientOp final : public Operator<Context> {
40  public:
41  ROIPoolGradientOp(const OperatorDef& def, Workspace* ws)
42  : Operator<Context>(def, ws),
43  pool_h_(OpArg<int64_t>("pool_h", 0)),
44  pool_w_(OpArg<int64_t>("pool_w", 0)),
45  spatial_scale_(OpArg<float>("spatial_scale", 1.f)) {}
47 
48  void RunImplFloat16();
49 
50  void RunOnDevice() override;
51  template <typename T> void RunImpl();
52 
53  protected:
55  int64_t pool_h_, pool_w_;
56 };
57 
58 } // namespace dragon
59 
60 #endif // DRAGON_OPERATORS_VISION_ROI_POOL_OP_H_
ROIPoolGradientOp(const OperatorDef &def, Workspace *ws)
Definition: roi_pool_op.h:41
void RunOnDevice() override
Implement the detailed execution.
Definition: roi_pool_op.cc:32
Definition: workspace.h:20
int64_t pool_w_
Definition: roi_pool_op.h:55
const OperatorDef & def() const
Return the stored def.
Definition: operator.h:114
Workspace * ws() const
Return the parent workspace.
Definition: operator.h:87
void RunImplFloat16()
Definition: roi_pool_op.cc:68
USE_OPERATOR_FUNCTIONS
Definition: roi_pool_op.h:46
Definition: roi_pool_op.h:39
float spatial_scale_
Definition: roi_pool_op.h:34
float spatial_scale_
Definition: roi_pool_op.h:54
void RunImpl()
Definition: roi_pool_op.cc:9
Definition: operator.h:149
#define OpArg
Definition: operator.h:235
int64_t pool_h_
Definition: roi_pool_op.h:35
int64_t pool_w_
Definition: roi_pool_op.h:35
USE_OPERATOR_FUNCTIONS
Definition: roi_pool_op.h:28
int64_t pool_h_
Definition: roi_pool_op.h:55
void RunOnDevice() override
Implement the detailed execution.
Definition: roi_pool_op.cc:111
ROIPoolOp(const OperatorDef &def, Workspace *ws)
Definition: roi_pool_op.h:23
Definition: roi_pool_op.h:21
Definition: common.h:41
void RunImpl()
Definition: roi_pool_op.cc:45