libtorch调用模型

 

这个也是调用模型,读取图片:

https://github.com/zacario-li/libtorch_cpp_mobilenetv2_5classes/blob/master/c%2B%2B/src/model_load.cpp

 

https://github.com/BIGBALLON/PyTorch-CPP/blob/master/prediction.cpp

 

 std::shared_ptr module =
      torch::jit::load(argv[1]);
  std::cout << "== Switch to GPU mode" << std::endl;
  // to GPU
  module->to(at::kCUDA);

  assert(module != nullptr);
  std::cout << "== ResNet50 loaded!\n";
  std::vector labels;
  if (LoadImageNetLabel(argv[2], labels)) {
    std::cout << "== Label loaded! Let's try it\n";
  } else {
    std::cerr << "Please check your label file path." << std::endl;
    return -1;
  }

  std::string file_name = "";
  cv::Mat image;
  while (true) {
    std::cout << "== Input image path: [enter Q to exit]" << std::endl;
    std::cin >> file_name;
    if (file_name == "Q") {
      break;
    }
    if (LoadImage(file_name, image)) {
      auto input_tensor = torch::from_blob(
          image.data, {1, kIMAGE_SIZE, kIMAGE_SIZE, kCHANNELS});
      input_tensor = input_tensor.permute({0, 3, 1, 2});
      input_tensor[0][0] = input_tensor[0][0].sub_(0.485).div_(0.229);
      input_tensor[0][1] = input_tensor[0][1].sub_(0.456).div_(0.224);
      input_tensor[0][2] = input_tensor[0][2].sub_(0.406).div_(0.225);

      // to GPU
      input_tensor = input_tensor.to(at::kCUDA);

      torch::Tensor out_tensor = module->forward({input_tensor}).toTensor();

 

你可能感兴趣的:(c++,torch)