()
| 127 | |
| 128 | |
| 129 | def main(): |
| 130 | common.add_help(description="Runs an MNIST network using a PyTorch model") |
| 131 | # Train the PyTorch model |
| 132 | mnist_model = model.MnistModel() |
| 133 | mnist_model.learn() |
| 134 | weights = mnist_model.get_weights() |
| 135 | # Do inference with TensorRT. |
| 136 | engine = build_engine(weights) |
| 137 | |
| 138 | # Build an engine, allocate buffers and create a stream. |
| 139 | # For more information on buffer allocation, refer to the introductory samples. |
| 140 | inputs, outputs, bindings, stream = common.allocate_buffers(engine) |
| 141 | context = engine.create_execution_context() |
| 142 | |
| 143 | case_num = load_random_test_case(mnist_model, pagelocked_buffer=inputs[0].host) |
| 144 | # For more information on performing inference, refer to the introductory samples. |
| 145 | # The common.do_inference function will return a list of outputs - we only have one in this case. |
| 146 | [output] = common.do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) |
| 147 | pred = np.argmax(output) |
| 148 | common.free_buffers(inputs, outputs, stream) |
| 149 | print("Test Case: " + str(case_num)) |
| 150 | print("Prediction: " + str(pred)) |
| 151 | |
| 152 | |
| 153 | if __name__ == "__main__": |
no test coverage detected