network_trt.cpp 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. // Copyright 2021 TIER IV, Inc.
  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 "lidar_centerpoint/network/network_trt.hpp"
  15. namespace centerpoint
  16. {
  17. bool VoxelEncoderTRT::setProfile(
  18. nvinfer1::IBuilder & builder, nvinfer1::INetworkDefinition & network,
  19. nvinfer1::IBuilderConfig & config)
  20. {
  21. auto profile = builder.createOptimizationProfile();
  22. auto in_name = network.getInput(0)->getName();
  23. auto in_dims = nvinfer1::Dims3(
  24. config_.max_voxel_size_, config_.max_point_in_voxel_size_, config_.encoder_in_feature_size_);
  25. profile->setDimensions(in_name, nvinfer1::OptProfileSelector::kMIN, in_dims);
  26. profile->setDimensions(in_name, nvinfer1::OptProfileSelector::kOPT, in_dims);
  27. profile->setDimensions(in_name, nvinfer1::OptProfileSelector::kMAX, in_dims);
  28. auto out_name = network.getOutput(0)->getName();
  29. auto out_dims = nvinfer1::Dims2(config_.max_voxel_size_, config_.encoder_out_feature_size_);
  30. profile->setDimensions(out_name, nvinfer1::OptProfileSelector::kMIN, out_dims);
  31. profile->setDimensions(out_name, nvinfer1::OptProfileSelector::kOPT, out_dims);
  32. profile->setDimensions(out_name, nvinfer1::OptProfileSelector::kMAX, out_dims);
  33. config.addOptimizationProfile(profile);
  34. return true;
  35. }
  36. HeadTRT::HeadTRT(
  37. const std::vector<std::size_t> & out_channel_sizes, const CenterPointConfig & config,
  38. const bool verbose)
  39. : TensorRTWrapper(config, verbose), out_channel_sizes_(out_channel_sizes)
  40. {
  41. }
  42. bool HeadTRT::setProfile(
  43. nvinfer1::IBuilder & builder, nvinfer1::INetworkDefinition & network,
  44. nvinfer1::IBuilderConfig & config)
  45. {
  46. auto profile = builder.createOptimizationProfile();
  47. auto in_name = network.getInput(0)->getName();
  48. auto in_dims = nvinfer1::Dims4(
  49. config_.batch_size_, config_.encoder_out_feature_size_, config_.grid_size_y_,
  50. config_.grid_size_x_);
  51. profile->setDimensions(in_name, nvinfer1::OptProfileSelector::kMIN, in_dims);
  52. profile->setDimensions(in_name, nvinfer1::OptProfileSelector::kOPT, in_dims);
  53. profile->setDimensions(in_name, nvinfer1::OptProfileSelector::kMAX, in_dims);
  54. for (std::size_t ci = 0; ci < out_channel_sizes_.size(); ci++) {
  55. auto out_name = network.getOutput(ci)->getName();
  56. auto out_dims = nvinfer1::Dims4(
  57. config_.batch_size_, out_channel_sizes_[ci], config_.down_grid_size_y_,
  58. config_.down_grid_size_x_);
  59. profile->setDimensions(out_name, nvinfer1::OptProfileSelector::kMIN, out_dims);
  60. profile->setDimensions(out_name, nvinfer1::OptProfileSelector::kOPT, out_dims);
  61. profile->setDimensions(out_name, nvinfer1::OptProfileSelector::kMAX, out_dims);
  62. }
  63. config.addOptimizationProfile(profile);
  64. return true;
  65. }
  66. } // namespace centerpoint