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
|
#pragma once
#include "caffe2/predictor/predictor.h"
#include "caffe2/proto/caffe2_pb.h"
#include <memory>
#include <string>
#include <vector>
namespace caffe2 {
namespace onnx {
class TORCH_API Caffe2BackendRep {
public:
void Run(
const caffe2::Predictor::TensorList& inputs,
caffe2::Predictor::TensorList* outputs);
void RunMap(
const caffe2::Predictor::TensorMap& inputs,
caffe2::Predictor::TensorList* outputs);
caffe2::NetDef& init_net() {
return init_net_;
}
caffe2::NetDef& pred_net() {
return pred_net_;
}
std::vector<std::string>& uninitialized_inputs() {
return uninitialized_inputs_;
}
const caffe2::NetDef& init_net() const {
return init_net_;
}
const caffe2::NetDef& pred_net() const {
return pred_net_;
}
const std::vector<std::string>& uninitialized_inputs() const {
return uninitialized_inputs_;
}
private:
void CheckInit();
caffe2::NetDef init_net_;
caffe2::NetDef pred_net_;
std::vector<std::string> uninitialized_inputs_;
std::unique_ptr<caffe2::Predictor> predictor_{nullptr};
};
} // namespace onnx
} // namespace caffe2
|
