Executes the network.
| 92 | |
| 93 | // Executes the network. |
| 94 | void Execute(nvinfer1::IExecutionContext* context, const float* input, |
| 95 | float* output) { |
| 96 | const nvinfer1::ICudaEngine& engine = context->getEngine(); |
| 97 | |
| 98 | // We have two bindings: input and output. |
| 99 | ASSERT_EQ(engine.getNbBindings(), 2); |
| 100 | const int input_index = engine.getBindingIndex(kInputTensor); |
| 101 | const int output_index = engine.getBindingIndex(kOutputTensor); |
| 102 | |
| 103 | // Create GPU buffers and a stream |
| 104 | void* buffers[2]; |
| 105 | ASSERT_EQ(0, cudaMalloc(&buffers[input_index], sizeof(float))); |
| 106 | ASSERT_EQ(0, cudaMalloc(&buffers[output_index], sizeof(float))); |
| 107 | cudaStream_t stream; |
| 108 | ASSERT_EQ(0, cudaStreamCreate(&stream)); |
| 109 | |
| 110 | // Copy the input to the GPU, execute the network, and copy the output back. |
| 111 | // |
| 112 | // Note that since the host buffer was not created as pinned memory, these |
| 113 | // async copies are turned into sync copies. So the following synchronization |
| 114 | // could be removed. |
| 115 | ASSERT_EQ(0, cudaMemcpyAsync(buffers[input_index], input, sizeof(float), |
| 116 | cudaMemcpyHostToDevice, stream)); |
| 117 | context->enqueue(1, buffers, stream, nullptr); |
| 118 | ASSERT_EQ(0, cudaMemcpyAsync(output, buffers[output_index], sizeof(float), |
| 119 | cudaMemcpyDeviceToHost, stream)); |
| 120 | cudaStreamSynchronize(stream); |
| 121 | |
| 122 | // Release the stream and the buffers |
| 123 | ASSERT_EQ(0, cudaFree(buffers[input_index])); |
| 124 | ASSERT_EQ(0, cudaFree(buffers[output_index])); |
| 125 | cudaStreamDestroy(stream); |
| 126 | } |
| 127 | |
| 128 | TEST(TensorrtTest, BasicFunctions) { |
| 129 | // Handle the case where the test is run on machine with no gpu available. |
no test coverage detected