[Question] How to correctly adapt the cpp predictor? about caffe2-tutorial HOT 2 CLOSED

/*******************************************************
 * Copyright (C) 2018-2019 bigballon <[email protected]>
 *
 * This file is a caffe2 C++ image classification test
 * by using pre-trained cifar10 model.
 *
 * Feel free to modify if you need.
 *******************************************************/
#include "caffe2/core/common.h"
#include "caffe2/core/init.h"
#include "caffe2/core/tensor.h"
#include "caffe2/core/workspace.h"
#include "caffe2/utils/proto_utils.h"

// feel free to define USE_GPU if you want to use gpu

// #define USE_GPU

#ifdef USE_GPU
#include "caffe2/core/context_gpu.h"
#endif

// headers for opencv
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include <iostream>
#include <map>
#include <string>

// define flags
CAFFE2_DEFINE_string(init_net, "./init_net.pb",
                     "The given path to the init protobuffer.");
CAFFE2_DEFINE_string(predict_net, "./predict_net.pb",
                     "The given path to the predict protobuffer.");
CAFFE2_DEFINE_string(file, "./image_file.jpg", "The image file.");

namespace caffe2 {

void loadImage(std::string file_name, float *imgArray) {
  auto image = cv::imread(file_name);  // CV_8UC3
  std::cout << "== image size: " << image.size() << " ==" << std::endl;

  // scale image to fit
  cv::Size scale(28, 28);
  cv::resize(image, image, scale);
  std::cout << "== simply resize: " << image.size() << " ==" << std::endl;

  // convert [unsigned int] to [float]
  image.convertTo(image, CV_32FC1);
  auto it = image.begin<float>();
  for (unsigned i = 0; it != image.end<float>(); it++, i++) {
    *(imgArray + i) = (*it);
  }
}

void run() {
  // define a caffe2 Workspace
  Workspace workSpace;

  // define initNet and predictNet
  NetDef initNet, predictNet;

  // read protobuf
  CAFFE_ENFORCE(ReadProtoFromFile(FLAGS_init_net, &initNet));
  CAFFE_ENFORCE(ReadProtoFromFile(FLAGS_predict_net, &predictNet));

  // set device type
#ifdef USE_GPU
  predictNet.mutable_device_option()->set_device_type(CUDA);
  initNet.mutable_device_option()->set_device_type(CUDA);
  std::cout << "== GPU processing selected "
            << " ==" << std::endl;
#else
  predictNet.mutable_device_option()->set_device_type(CPU);
  initNet.mutable_device_option()->set_device_type(CPU);

  for (int i = 0; i < predictNet.op_size(); ++i) {
    predictNet.mutable_op(i)->mutable_device_option()->set_device_type(CPU);
  }
  for (int i = 0; i < initNet.op_size(); ++i) {
    initNet.mutable_op(i)->mutable_device_option()->set_device_type(CPU);
  }
#endif

  // load network
  CAFFE_ENFORCE(workSpace.RunNetOnce(initNet));
  CAFFE_ENFORCE(workSpace.CreateNet(predictNet));
  std::cout << "== network loaded "
            << " ==" << std::endl;

  // load image from file, then convert it to float array.
  float imgArray[1 * 28 * 28];
  loadImage(FLAGS_file, imgArray);

  // define a Tensor which is used to stone input data
  std::cout << "== Tensor "
            << " ==" << std::endl;
  TensorCPU input;
  input.Resize(std::vector<TIndex>({1, 1, 28, 28}));
  input.ShareExternalPointer(imgArray);

  // get "data" blob
#ifdef USE_GPU
  auto data = workSpace.GetBlob("data")->GetMutable<TensorCUDA>();
#else
  auto data = workSpace.GetBlob("data")->GetMutable<TensorCPU>();
#endif
  std::cout << "== Blob got "
            << " ==" << std::endl;

  // copy from input data
  data->CopyFrom(input);
  std::cout << "== Data copied "
            << " ==" << std::endl;

  // forward
  workSpace.RunNet(predictNet.name());
  std::cout << "== Net was run "
            << " ==" << std::endl;

  // get predictions blob and show the results
  std::vector<std::string> labelName = {"0", "1", "2", "3", "4",
                                        "5", "6", "7", "8", "9"};

#ifdef USE_GPU
  auto predictions =
      TensorCPU(workSpace.GetBlob("predictions")->Get<TensorCUDA>());
#else
  auto predictions = workSpace.GetBlob("predictions")->Get<TensorCPU>();
#endif
  std::cout << "== predictions got "
            << " ==" << std::endl;

  std::vector<float> probs(predictions.data<float>(),
                           predictions.data<float>() + predictions.size());

  auto max = std::max_element(probs.begin(), probs.end());
  auto index = std::distance(probs.begin(), max);
  std::cout << "== predicted label: " << labelName[index]
            << " ==\n== with probability: " << (*max * 100)
            << "% ==" << std::endl;
}

}  // namespace caffe2

// main function
int main(int argc, char **argv) {
  caffe2::GlobalInit(&argc, &argv);
  caffe2::run();
  google::protobuf::ShutdownProtobufLibrary();
  return 0;
}