-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.cpp
70 lines (58 loc) · 2.26 KB
/
main.cpp
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
#include "tfutil.h"
int loadTFModelFile(const char* fileName, std::vector<std::vector<std::string>>& net, int inputDim[4], std::string outputFolder)
{
//verify the version of protobuf library.
GOOGLE_PROTOBUF_VERIFY_VERSION;
//read the TF model and initializing graph
std:: cout<<"Reading the binary file from : "<< fileName << std::endl;
std::ifstream input(fileName, std::ios::in | std::ios::binary);
tensorflow::GraphDef graph;
if (!input.is_open())
{
std::cout << "cant open :" << fileName << std::endl;
return -1;
}
bool isSuccess = graph.ParseFromIstream(&input);
if(isSuccess) {
std::cout << "TF Model Read Successful" << std::endl;
int node_param_size = graph.node_size();
if(node_param_size > 0) {
TFUtil tf_util = TFUtil(&graph);
tf_util.PrintModel();
//Write GDF
const char * tensorType = "VX_TYPE_FLOAT32";
const char * convertPolicy = "VX_CONVERT_POLICY_SATURATE";
const char * roundPolicy = "VX_ROUND_POLICY_TO_NEAREST_EVEN";
int fixedPointPosition = 0;
bool isVirtualEnabled = true;
bool bFuseScaleWithBatchNorm = true;
tf_util.GenerateGDF(inputDim, outputFolder, tensorType, convertPolicy, roundPolicy, fixedPointPosition, isVirtualEnabled, bFuseScaleWithBatchNorm);
}
else {
std::cerr << "ERROR: [Unsupported caffemodel] please upgrade this caffemodel, currently uses deprecated V1LayerParameters." << std::endl;
return -1;
}
}
else {
std::cerr << "Tensorflow Model Read Failed" << std::endl;
}
input.close();
return 0;
}
int main(int argc, char* argv[])
{
const char *fileName = "Models/MNIST/MNIST.pb";
std::vector<std::vector<std::string>> net;
int inputDim[4] = { 1, 1, 28, 28 }; //[Batch Size, Channel, Height, Width]
std::string outputFolder = ".";
// make sure that weights and bias folder are created
std::string dir = outputFolder + "/weights";
mkdir(dir.c_str(), 0777);
dir = outputFolder + "/bias";
mkdir(dir.c_str(), 0777);
// load TF model
if(loadTFModelFile(fileName, net, inputDim, outputFolder) < 0) {
return -1;
}
return 0;
}