| 126 | } |
| 127 | |
| 128 | TEST(TensorrtTest, BasicFunctions) { |
| 129 | // Handle the case where the test is run on machine with no gpu available. |
| 130 | if (CHECK_NOTNULL(GPUMachineManager())->VisibleDeviceCount() <= 0) { |
| 131 | LOG(WARNING) << "No gpu device available, probably not being run on a gpu " |
| 132 | "machine. Skipping..."; |
| 133 | return; |
| 134 | } |
| 135 | |
| 136 | // Create a serialized engine |
| 137 | TrtUniquePtrType<nvinfer1::IHostMemory> model = CreateSerializedEngine(); |
| 138 | // Use the model to create an engine and then an execution context. |
| 139 | Logger& logger = *Logger::GetLogger(); |
| 140 | TrtUniquePtrType<nvinfer1::IRuntime> runtime( |
| 141 | nvinfer1::createInferRuntime(logger)); |
| 142 | TrtUniquePtrType<nvinfer1::ICudaEngine> engine( |
| 143 | runtime->deserializeCudaEngine(model->data(), model->size(), nullptr)); |
| 144 | TrtUniquePtrType<nvinfer1::IExecutionContext> context( |
| 145 | engine->createExecutionContext()); |
| 146 | |
| 147 | // Execute the network. |
| 148 | float input = 1234; |
| 149 | float output; |
| 150 | Execute(context.get(), &input, &output); |
| 151 | EXPECT_EQ(output, input * 2 + 3); |
| 152 | } |
| 153 | |
| 154 | } // namespace tensorrt |
| 155 | } // namespace tensorflow |
nothing calls this directly
no test coverage detected