multi_gpu_model_infer.cpp 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. // Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. #include <gflags/gflags.h>
  15. #include <omp.h>
  16. #include <memory>
  17. #include <string>
  18. #include <fstream>
  19. #include "model_deploy/common/include/multi_gpu_model.h"
  20. DEFINE_string(model_filename, "", "Path of det inference model");
  21. DEFINE_string(params_filename, "", "Path of det inference params");
  22. DEFINE_string(cfg_file, "", "Path of yaml file");
  23. DEFINE_string(model_type, "", "model type");
  24. DEFINE_string(image_list, "", "Path of test image file");
  25. DEFINE_string(gpu_id, "0", "GPU card id, example: 0,2,3");
  26. DEFINE_int32(batch_size, 1, "Batch size of infering");
  27. DEFINE_int32(thread_num, 1, "thread num of preprocessing");
  28. int main(int argc, char** argv) {
  29. google::ParseCommandLineFlags(&argc, &argv, true);
  30. std::vector<int> gpu_ids;
  31. std::stringstream gpu_ids_str(FLAGS_gpu_id);
  32. std::string temp;
  33. while (getline(gpu_ids_str, temp, ',')) {
  34. gpu_ids.push_back(std::stoi(temp));
  35. }
  36. for (auto gpu_id : gpu_ids) {
  37. std::cout << "gpu_id:" << gpu_id << std::endl;
  38. }
  39. std::cout << "start create model" << std::endl;
  40. // create model
  41. PaddleDeploy::MultiGPUModel model;
  42. if (!model.Init(FLAGS_model_type, FLAGS_cfg_file, gpu_ids.size())) {
  43. return -1;
  44. }
  45. // engine init
  46. PaddleDeploy::PaddleEngineConfig engine_config;
  47. engine_config.model_filename = FLAGS_model_filename;
  48. engine_config.params_filename = FLAGS_params_filename;
  49. engine_config.use_gpu = true;
  50. engine_config.max_batch_size = FLAGS_batch_size;
  51. if (!model.PaddleEngineInit(engine_config, gpu_ids)) {
  52. return -1;
  53. }
  54. // Mini-batch
  55. if (FLAGS_image_list == "") {
  56. std::cerr << "image_list should be defined" << std::endl;
  57. return -1;
  58. }
  59. std::vector<std::string> image_paths;
  60. std::ifstream inf(FLAGS_image_list);
  61. if (!inf) {
  62. std::cerr << "Fail to open file " << FLAGS_image_list << std::endl;
  63. return -1;
  64. }
  65. std::string image_path;
  66. while (getline(inf, image_path)) {
  67. image_paths.push_back(image_path);
  68. }
  69. std::cout << "start model predict " << image_paths.size() << std::endl;
  70. // infer
  71. std::vector<PaddleDeploy::Result> results;
  72. for (int i = 0; i < image_paths.size(); i += FLAGS_batch_size) {
  73. // Read image
  74. int im_vec_size =
  75. std::min(static_cast<int>(image_paths.size()), i + FLAGS_batch_size);
  76. std::vector<cv::Mat> im_vec(im_vec_size - i);
  77. #pragma omp parallel for num_threads(im_vec_size - i)
  78. for (int j = i; j < im_vec_size; ++j) {
  79. im_vec[j - i] = std::move(cv::imread(image_paths[j], 1));
  80. }
  81. model.Predict(im_vec, &results, FLAGS_thread_num);
  82. std::cout << i / FLAGS_batch_size << " group" << std::endl;
  83. for (auto j = 0; j < results.size(); ++j) {
  84. std::cout << "Result for sample " << j << std::endl;
  85. std::cout << results[j] << std::endl;
  86. }
  87. }
  88. return 0;
  89. }