CHNN/transform_caffe_model.cpp at master · camel007/CHNN · GitHub

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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
#include <caffe/caffe.hpp>
#ifdef USE_OPENCV
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#endif  // USE_OPENCV
#include <algorithm>
#include <iosfwd>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include <cstdio>

#ifdef USE_OPENCV
using namespace caffe;  // NOLINT(build/namespaces)
using std::string;

/* Pair (label, confidence) representing a prediction. */
typedef std::pair<string, float> Prediction;

class Classifier {
 public:
  Classifier(const string& model_file,
             const string& trained_file);

  std::vector<Prediction> Classify(const cv::Mat& img, int N = 5);

  void saveCaffeModel();

 private:
  void SetMean(const string& mean_file);

  std::vector<float> Predict(const cv::Mat& img);

  void WrapInputLayer(std::vector<cv::Mat>* input_channels);

  void Preprocess(const cv::Mat& img,
                  std::vector<cv::Mat>* input_channels);

 private:
  shared_ptr<Net<float> > net_;
  cv::Size input_geometry_;
  int num_channels_;
  cv::Mat mean_;
  std::vector<string> labels_;
};


void Classifier::saveCaffeModel()
{
    vector<shared_ptr<Layer<float> > > layers = net_->layers();
    int layer_sz = layers.size();

    std::vector<std::string> layer_names = net_->layer_names();

    std::FILE *fout = std::fopen("vgg_16_face.bin", "w");
    if(!fout) return;
    for(int i = 0; i < layer_sz - 4; ++i)
    {
        boost::shared_ptr<caffe::Layer<float> > lay = layers[i];
        std::string name = lay->type();
        std::vector<boost::shared_ptr<caffe::Blob<float> > >& blobs = lay->blobs();

        std::string layer_name = layer_names[i];

        if(name == "Convolution")
        {
            boost::shared_ptr<caffe::Blob<float> > weight = blobs[0];
            boost::shared_ptr<caffe::Blob<float> > bias = blobs[1];

            int pad = 1;
            //if(layer_name == "face-conv6")
            //   pad = 0;
            int stride = 1;
            std::fwrite(&pad, sizeof(int), 1, fout);
            std::fwrite(&stride, sizeof(int), 1, fout);

            //weight
            std::vector<int> shape = weight->shape();
            std::fwrite(&shape[0], sizeof(int), 4, fout);
            std::fwrite(weight->cpu_data(), sizeof(float), weight->count(), fout);

            //bias
            shape = bias->shape();
            std::fwrite(&shape[0], sizeof(float), 1, fout);
            std::fwrite(bias->cpu_data(), sizeof(float), bias->count(), fout);
        }
        else if(name == "Pooling")
        {
            int pad = 0;
            int stride = 2;
            int kernel_size = 2;
            std::fwrite(&pad, sizeof(int), 1, fout);
            std::fwrite(&stride, sizeof(int), 1, fout);
            std::fwrite(&kernel_size, sizeof(int), 1, fout);
        }
        else if(name == "InnerProduct")
        {
            boost::shared_ptr<caffe::Blob<float> > weight = blobs[0];
            boost::shared_ptr<caffe::Blob<float> > bias = blobs[1];

            //weight
            std::vector<int> shape = weight->shape();
            std::fwrite(&shape[0], sizeof(int), 2, fout);
            std::fwrite(weight->cpu_data(), sizeof(float), weight->count(), fout);

            //bias
            shape = bias->shape();
            std::fwrite(&shape[0], sizeof(int), 1, fout);
            std::fwrite(bias->cpu_data(), sizeof(float), bias->count(), fout);
        }
    }

    std::fclose(fout);
}

int main() {

  ::google::InitGoogleLogging("");

  string model_file   = "/home/rpk/caffe-master/models/vgg_face/VGG_FACE_deploy.prototxt";
  string trained_file = "/home/rpk/caffe-master/models/vgg_face/VGG_FACE.caffemodel";
  Classifier classifier(model_file, trained_file);
  classifier.saveCaffeModel();

}
#else
int main(int argc, char** argv) {
  LOG(FATAL) << "This example requires OpenCV; compile with USE_OPENCV.";
}
#endif  // USE_OPENCV