1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
|
#ifndef CAFFE2_OPERATORS_IF_OP_H_
#define CAFFE2_OPERATORS_IF_OP_H_
#include "caffe2/core/context.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/operator.h"
namespace caffe2 {
template <class Context>
class IfOp final : public Operator<Context> {
public:
IfOp(const OperatorDef& operator_def, Workspace* ws)
: Operator<Context>(operator_def, ws) {
CAFFE_ENFORCE(
this->template HasSingleArgumentOfType<NetDef>("then_net"),
"then_net must be specified in If operator");
auto then_net_def =
this->template GetSingleArgument<NetDef>("then_net", NetDef());
then_net_ = CreateNet(then_net_def, ws);
CAFFE_ENFORCE(then_net_, "Failed to initialize then subnet");
if (this->template HasSingleArgumentOfType<NetDef>("else_net")) {
auto else_net_def =
this->template GetSingleArgument<NetDef>("else_net", NetDef());
else_net_ = CreateNet(else_net_def, ws);
CAFFE_ENFORCE(else_net_, "Failed to initialize else subnet");
}
}
USE_OPERATOR_CONTEXT_FUNCTIONS;
bool RunOnDevice() override {
CAFFE_ENFORCE(
this->InputIsTensorType(0, Context::GetDeviceType()),
"Invalid condition in If operator: tensor expected");
const auto& condition = Input(0);
CAFFE_ENFORCE_EQ(
condition.numel(),
1,
"Invalid condition tensor in If operator: single value expected");
auto conditionValue = *condition.template data<bool>();
if (conditionValue) {
return then_net_->Run();
} else if (else_net_) {
return else_net_->Run();
}
return true;
}
private:
std::unique_ptr<NetBase> then_net_;
std::unique_ptr<NetBase> else_net_;
};
} // namespace caffe2
#endif // CAFFE2_OPERATORS_IF_OP_H_
|
