()
| 78 | |
| 79 | |
| 80 | def main(): |
| 81 | # Set the data path to the directory that contains the trained models and test images for inference. |
| 82 | _, data_files = common.find_sample_data( |
| 83 | description="Runs a ResNet50 network with a TensorRT inference engine.", |
| 84 | subfolder="resnet50", |
| 85 | find_files=[ |
| 86 | "binoculars.jpeg", |
| 87 | "reflex_camera.jpeg", |
| 88 | "tabby_tiger_cat.jpg", |
| 89 | ModelData.MODEL_PATH, |
| 90 | "class_labels.txt", |
| 91 | ], |
| 92 | ) |
| 93 | # Get test images, models and labels. |
| 94 | test_images = data_files[0:3] |
| 95 | onnx_model_file, labels_file = data_files[3:] |
| 96 | labels = open(labels_file, "r").read().split("\n") |
| 97 | |
| 98 | # Build a TensorRT engine. |
| 99 | engine = build_engine_onnx(onnx_model_file) |
| 100 | # Inference is the same regardless of which parser is used to build the engine, since the model architecture is the same. |
| 101 | # Allocate buffers and create a CUDA stream. |
| 102 | inputs, outputs, bindings, stream = common.allocate_buffers(engine) |
| 103 | # Contexts are used to perform inference. |
| 104 | context = engine.create_execution_context() |
| 105 | |
| 106 | # Load a normalized test case into the host input page-locked buffer. |
| 107 | test_image = random.choice(test_images) |
| 108 | test_case = load_normalized_test_case(test_image, inputs[0].host) |
| 109 | # Run the engine. The output will be a 1D tensor of length 1000, where each value represents the |
| 110 | # probability that the image corresponds to that label |
| 111 | trt_outputs = common.do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) |
| 112 | # We use the highest probability as our prediction. Its index corresponds to the predicted label. |
| 113 | pred = labels[np.argmax(trt_outputs[0])] |
| 114 | common.free_buffers(inputs, outputs, stream) |
| 115 | if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]: |
| 116 | print("Correctly recognized " + test_case + " as " + pred) |
| 117 | else: |
| 118 | print("Incorrectly recognized " + test_case + " as " + pred) |
| 119 | |
| 120 | |
| 121 | if __name__ == "__main__": |
no test coverage detected