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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
|
#include "caffe2/core/operator.h"
#include "caffe2/core/tensor.h"
#include "caffe2/core/types.h"
#include "caffe2/utils/proto_utils.h"
#include "NeuralNetworks.h"
#include "dlnnapi.h"
namespace caffe2 {
class NNApi {
public:
using TensorVector = std::vector<TensorCPU*>;
// Three different modes:
// ANEURALNETWORKS_PREFER_LOW_POWER
// ANEURALNETWORKS_PREFER_FAST_SINGLE_ANSWER
// ANEURALNETWORKS_PREFER_SUSTAINED_SPEED
NNApi(
const NetDef& init_net,
const NetDef& run_net,
Workspace* ws = nullptr,
const PreferenceCode pref = ANEURALNETWORKS_PREFER_SUSTAINED_SPEED)
: preference_(pref),
run_net_(run_net),
ws_(ws) {
if (!loadNNApiLibrary()) {
CAFFE_THROW("NNApi is not supported");
}
CAFFE_ENFORCE(ws_.RunNetOnce(init_net));
}
~NNApi();
bool loadNNApiLibrary();
bool run(const TensorVector& inputs, TensorVector* outputs);
private:
dlnnapi libnnapi_;
ANeuralNetworksModel* model_{nullptr};
ANeuralNetworksCompilation* compilation_{nullptr};
ANeuralNetworksExecution* run_{nullptr};
ANeuralNetworksEvent* run_end_{nullptr};
PreferenceCode preference_;
NetDef run_net_;
Workspace ws_;
OperandCode tensor_type_;
uint32_t operand_idx{0};
std::unordered_map<std::string, uint32_t> operand_map_;
// dimensions for the tensors
std::unordered_map<std::string, std::vector<uint32_t>> tensor_dims_;
// mapping of the operator name "Conv" to OperatorType CONV
enum OperatorType {
AVERAGEPOOL,
CONV,
MAXPOOL,
RELU,
SOFTMAX,
};
std::unordered_map<std::string, OperatorType> operator_map_{
{"AveragePool", AVERAGEPOOL},
{"Conv", CONV},
{"MaxPool", MAXPOOL},
{"Relu", RELU},
{"Softmax", SOFTMAX}};
struct ConvPoolArgs {
int kernel_h{0};
int kernel_w{0};
int stride_x{0};
int stride_y{0};
int pad_t{0};
int pad_l{0};
int pad_b{0};
int pad_r{0};
};
void getConvPoolArgs(const ArgumentHelper& helper, ConvPoolArgs& args);
uint32_t addScalarOperand(int32_t val);
uint32_t addFloatOperand(float val);
uint32_t addTensorOperand(
const std::string& blob,
OperandCode type,
std::vector<uint32_t>& dims,
float scale = 1.0,
int32_t zero_point = 0);
// lazily initialize model_ in run()
void init(const TensorVector& inputs, TensorVector* outputs);
void addConv(const OperatorDef& op, bool fuse_relu = false);
void addPooling(
const OperatorDef& op,
OperationCode op_code,
bool fuse_relu = false);
void addRelu(const OperatorDef& op);
void addSoftmax(const OperatorDef& op);
};
} // namespace caffe2
|
