libtorch呼び出しモデル

1670 ワード


これも呼び出しモデルで、画像を読み込みます.
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<:jit::script::module> 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!
"; std::vector<:string> labels; if (LoadImageNetLabel(argv[2], labels)) { std::cout << "== Label loaded! Let's try it
"; } 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();