Caffe2 - C++ API
A deep learning, cross platform ML framework
prelu_op.h
1 
17 #pragma once
18 
19 #include "caffe2/core/context.h"
20 #include "caffe2/core/logging.h"
21 #include "caffe2/core/operator.h"
22 #include "caffe2/utils/math.h"
23 
24 namespace caffe2 {
25 
26 template <typename T, class Context>
27 class PReluOp final : public Operator<Context> {
28  public:
29  PReluOp(const OperatorDef& operator_def, Workspace* ws)
30  : Operator<Context>(operator_def, ws),
31  order_(StringToStorageOrder(
32  OperatorBase::GetSingleArgument<string>("order", "NCHW"))) {}
33 
34  USE_OPERATOR_CONTEXT_FUNCTIONS;
35 
36  bool RunOnDevice() override;
37 
38  protected:
39  StorageOrder order_;
40 };
41 
42 template <typename T, class Context>
43 class PReluGradientOp final : public Operator<Context> {
44  public:
45  PReluGradientOp(const OperatorDef& operator_def, Workspace* ws)
46  : Operator<Context>(operator_def, ws),
47  order_(StringToStorageOrder(
48  OperatorBase::GetSingleArgument<string>("order", "NCHW"))) {}
49  USE_OPERATOR_CONTEXT_FUNCTIONS;
50 
51  bool RunOnDevice() override;
52 
53  protected:
54  StorageOrder order_;
55 };
56 
57 } // namespace caffe2
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.