\brief Runs the TensorRT inference. \details Allocate input and output memory, and executes the engine.
| 86 | //! \details Allocate input and output memory, and executes the engine. |
| 87 | //! |
| 88 | bool SampleSegmentation::infer(const std::string& input_filename, int32_t width, int32_t height, const std::string& output_filename) |
| 89 | { |
| 90 | auto context = util::UniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext()); |
| 91 | if (!context) |
| 92 | { |
| 93 | return false; |
| 94 | } |
| 95 | |
| 96 | auto input_idx = mEngine->getBindingIndex("input"); |
| 97 | if (input_idx == -1) |
| 98 | { |
| 99 | return false; |
| 100 | } |
| 101 | assert(mEngine->getBindingDataType(input_idx) == nvinfer1::DataType::kFLOAT); |
| 102 | auto input_dims = nvinfer1::Dims4{1, 3 /* channels */, height, width}; |
| 103 | context->setBindingDimensions(input_idx, input_dims); |
| 104 | auto input_size = util::getMemorySize(input_dims, sizeof(float)); |
| 105 | |
| 106 | auto output_idx = mEngine->getBindingIndex("output"); |
| 107 | if (output_idx == -1) |
| 108 | { |
| 109 | return false; |
| 110 | } |
| 111 | assert(mEngine->getBindingDataType(output_idx) == nvinfer1::DataType::kINT32); |
| 112 | auto output_dims = context->getBindingDimensions(output_idx); |
| 113 | auto output_size = util::getMemorySize(output_dims, sizeof(int32_t)); |
| 114 | |
| 115 | // Allocate CUDA memory for input and output bindings |
| 116 | void* input_mem{nullptr}; |
| 117 | if (cudaMalloc(&input_mem, input_size) != cudaSuccess) |
| 118 | { |
| 119 | gLogError << "ERROR: input cuda memory allocation failed, size = " << input_size << " bytes" << std::endl; |
| 120 | return false; |
| 121 | } |
| 122 | void* output_mem{nullptr}; |
| 123 | if (cudaMalloc(&output_mem, output_size) != cudaSuccess) |
| 124 | { |
| 125 | gLogError << "ERROR: output cuda memory allocation failed, size = " << output_size << " bytes" << std::endl; |
| 126 | return false; |
| 127 | } |
| 128 | |
| 129 | // Read image data from file and mean-normalize it |
| 130 | const std::vector<float> mean{0.485f, 0.456f, 0.406f}; |
| 131 | const std::vector<float> stddev{0.229f, 0.224f, 0.225f}; |
| 132 | auto input_image{util::RGBImageReader(input_filename, input_dims, mean, stddev)}; |
| 133 | input_image.read(); |
| 134 | auto input_buffer = input_image.process(); |
| 135 | cudaStream_t stream; |
| 136 | if (cudaStreamCreate(&stream) != cudaSuccess) |
| 137 | { |
| 138 | gLogError << "ERROR: cuda stream creation failed." << std::endl; |
| 139 | return false; |
| 140 | } |
| 141 | |
| 142 | // Copy image data to input binding memory |
| 143 | if (cudaMemcpyAsync(input_mem, input_buffer.get(), input_size, cudaMemcpyHostToDevice, stream) != cudaSuccess) |
| 144 | { |
| 145 | gLogError << "ERROR: CUDA memory copy of input failed, size = " << input_size << " bytes" << std::endl; |
no test coverage detected