Caffe2——C++ 預測(predict)Demo

由於最近入坑Caffe2，它最近還一直在更新，因此坑比較多，官方也只給出了python的demo，C++的暫時還找不到，有也只有一個簡單版的，不夠用，因此就總結了一下，結合網上和本身的實踐，整理了一下代碼。

#include "caffe2/core/flags.h"
#include "caffe2/core/init.h"
#include "caffe2/core/predictor.h"
#include "caffe2/utils/proto_utils.h"
#include <opencv2/opencv.hpp>
#include <ctime>

namespace caffe2 {

    std::unique_ptr<Blob> randomTensor(
            const std::vector<TIndex>& dims,
            CPUContext* ctx
            ){
        auto blob = make_unique<Blob>();
        auto* t = blob->GetMutable<TensorCPU>();
        t->Resize(dims);
        math::RandUniform<float, CPUContext>(
                t->size(), -1.0, 1.0, t->template mutable_data<float>(), ctx
                );
        return blob;
    }

void run() {
  // 定義初始化網絡結構與權重值
  caffe2::NetDef init_net, predict_net;
  DeviceOption op;
  op.set_random_seed(1701);

  std::unique_ptr<CPUContext> ctx_;
  ctx_ = caffe2::make_unique<CPUContext>(op);

  // 讀入網絡結構文件
  ReadProtoFromFile("squeezenet/exec_net.pb", &init_net);
  ReadProtoFromFile("squeezenet/predict_net.pb", &predict_net);

  // Can be large due to constant fills
  VLOG(1) << "Init net: " << ProtoDebugString(init_net);
  LOG(INFO) << "Predict net: " << ProtoDebugString(predict_net);
  auto predictor = caffe2::make_unique<Predictor>(init_net, predict_net);
  LOG(INFO) << "Checking that a null forward-pass works";

  // 用opencv的方式讀入文件
  cv::Mat bgr_img = cv::imread("cat.jpg", -1);

  int height = bgr_img.rows;
  int width = bgr_img.cols;

  // 輸入圖像大小
  const int predHeight = 256;
  const int predWidth = 256;
  const int crops = 1;      // crops等於1表示batch的數量爲1
  const int channels = 3;   // 通道數爲3，表示BGR，爲1表示灰度圖
  const int size = predHeight * predWidth;
  const float hscale = ((float)height) / predHeight; // 計算縮放比例
  const float wscale = ((float)width) / predWidth;
  const float scale = std::min(hscale, wscale);
  // 初始化網絡的輸入，由於可能要作batch操做，因此分配一段連續的存儲空間
  std::vector<float> inputPlanar(crops * channels * predHeight * predWidth);

  std::cout << "before resizing, bgr_img.cols=" << bgr_img.cols << ", bgr_img.rows=" << bgr_img.rows << std::endl;
  // resize成想要的輸入大小
  cv::Size dsize = cv::Size(bgr_img.cols / wscale, bgr_img.rows / hscale);
  cv::resize(bgr_img, bgr_img, dsize);
  std::cout << "after resizing, bgr_img.cols=" << bgr_img.cols << ", bgr_img.rows=" << bgr_img.rows << std::endl;
  // Scale down the input to a reasonable predictor size.
  // 這裏是將圖像複製到連續的存儲空間內，用於網絡的輸入，由於是BGR三通道，因此有三個賦值
  // 注意imread讀入的圖像格式是unsigned char，若是你的網絡輸入要求是float的話，下面的操做就不對了。
  for (auto i=0; i<predHeight; i++) {
      //printf("+\n");
      for (auto j=0; j<predWidth; j++) {
          inputPlanar[i * predWidth + j + 0*size] = (float)bgr_img.data[(i*predWidth + j) * 3 + 0];
          inputPlanar[i * predWidth + j + 1*size] = (float)bgr_img.data[(i*predWidth + j) * 3 + 1];
          inputPlanar[i * predWidth + j + 2*size] = (float)bgr_img.data[(i*predWidth + j) * 3 + 2];
      }
  }
  // 輸入是float格式
  //for (auto i = 0; i < predHeight; i++) {
    // 模版的輸入格式是float
  //  const float* inData = bgr_img.ptr<float>(i);
  //  for (auto j = 0; j < predWidth; j++) {
  //      inputPlanar[i * predWidth + j + 0 * size] = (float)((inData[j]) * 3 ＋ 0);
  //      inputPlanar[i * predWidth + j + 1 * size] = (float)((inData[j]) * 3 ＋ 1);
  //      inputPlanar[i * predWidth + j + 2 * size] = (float)((inData[j]) * 3 ＋ 2);
  //  }
  //}


  //typedef Tensor<CPUContext> TensorCPU;
  // input就是網絡的輸入，因此把以前準備好的數據賦值給input就能夠了
  caffe2::TensorCPU input;
  input.Resize(std::vector<int>({crops, channels, predHeight, predWidth}));
  input.ShareExternalPointer(inputPlanar.data());

  //Predictor::TensorVector inputVec{inputData->template GetMutable<TensorCPU>()};
  Predictor::TensorVector inputVec{&input};

  Predictor::TensorVector outputVec;
  //predictor->run(inputVec, &outputVec);
  //CAFFE_ENFORCE_GT(outputVec.size(), 0);

  std::clock_t begin = clock();  //begin time of inference
  // 預測
  predictor->run(inputVec, &outputVec);

  //std::cout << "CAFFE2_LOG_THRESHOLD=" << CAFFE2_LOG_THRESHOLD << std::endl;
  //std::cout << "init_net.name()" << init_net.name();

  std::clock_t end = clock();
  double elapsed_secs = double(end-begin) / CLOCKS_PER_SEC;

  std::cout << "inference takes " << elapsed_secs << std::endl;

  float max_value = 0;
  int best_match_index = -1;
  // 迭代輸出結果，output的大小就是網絡輸出的大小
  for(auto output : outputVec) {
      for(auto i=0; i<output->size(); ++i){
          // val對應的就是每一類的機率值
          float val = output->template data<float>()[i];
          if(val>0.001){
              printf("%i: %s : %f\n", i, imagenet_classes[i], val);
              if(val>max_value) {
                  max_value = val;
                  best_match_index = i;
              }
          }
      }
  }
  // 這裏是用imagenet數據集爲例
  std::cout << "predicted result is:" << imagenet_classes[best_match_index] << ", with confidence of " << max_value << std::endl;

}
}

int main(int argc, char** argv) {
  caffe2::GlobalInit(&argc, &argv);
  caffe2::run();
  // This is to allow us to use memory leak checks.
  google::protobuf::ShutdownProtobufLibrary();
  return 0;
}