-
Notifications
You must be signed in to change notification settings - Fork 24
Expand file tree
/
Copy pathCaffe2Net.cpp
More file actions
62 lines (57 loc) · 1.58 KB
/
Copy pathCaffe2Net.cpp
File metadata and controls
62 lines (57 loc) · 1.58 KB
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
#include <caffe2/core/context_gpu.h>
#include "Caffe2Net.h"
Caffe2Net::Caffe2Net(string initNet,string predictNet)
:workspace(nullptr)
{
#ifdef WITH_CUDA
DeviceOption option;
option.set_device_type(CUDA);
new CUDAContext(option);
#endif
//载入部署模型
NetDef init_net_def, predict_net_def;
CAFFE_ENFORCE(ReadProtoFromFile(initNet, &init_net_def));
CAFFE_ENFORCE(ReadProtoFromFile(predictNet, &predict_net_def));
#ifdef WITH_CUDA
init_net_def.mutable_device_option()->set_device_type(CUDA);
predict_net_def.mutable_device_option()->set_device_type(CUDA);
#else
init_net_def.mutable_device_option()->set_device_type(CPU);
predict_net_def.mutable_device_option()->set_device_type(CPU);
#endif
//网络初始化
workspace.RunNetOnce(init_net_def);
//创建判别器
predict_net = CreateNet(predict_net_def,&workspace);
}
Caffe2Net::~Caffe2Net()
{
}
vector<float> Caffe2Net::predict(Mat img)
{
//create input blob
#ifdef WITH_CUDA
TensorCUDA input = TensorCUDA(preProcess(img));
auto tensor = workspace.CreateBlob("data")->GetMutable<TensorCUDA>();
#else
TensorCPU input = preProcess(img);
auto tensor = workspace.CreateBlob("data")->GetMutable<TensorCPU>();
#endif
tensor->ResizeLike(input);
tensor->ShareData(input);
//predict
predict_net->Run();
//get output blob
#ifdef WITH_CUDA
TensorCPU output = TensorCPU(workspace.GetBlob("softmax")->Get<TensorCUDA>());
#else
TensorCPU output = TensorCPU(workspace.GetBlob("softmax")->Get<TensorCPU>());
#endif
return postProcess(output);
}
TensorCPU Caffe2Net::preProcess(Mat img)
{
}
vector<float> Caffe2Net::postProcess(TensorCPU output)
{
}