(args, cfg)
| 89 | |
| 90 | |
| 91 | def check(args, cfg): |
| 92 | from mmdeploy.backend.tensorrt.wrapper import TRTWrapper |
| 93 | |
| 94 | model = get_model(args, cfg).cuda() |
| 95 | model.eval() |
| 96 | trt_model = TRTWrapper(f'{args.model_name}.engine', |
| 97 | ['output']) |
| 98 | |
| 99 | x = torch.randn(1, 3, args.size, args.size).cuda() |
| 100 | |
| 101 | torch_out = model(x) |
| 102 | trt_out = trt_model(dict(input=x))['output'] |
| 103 | |
| 104 | print('torch out shape:', torch_out.shape) |
| 105 | print('trt out shape:', trt_out.shape) |
| 106 | |
| 107 | print('max delta:', (torch_out - trt_out).abs().max()) |
| 108 | print('mean delta:', (torch_out - trt_out).abs().mean()) |
| 109 | |
| 110 | speed_test(model, x) |
| 111 | speed_test(trt_model, dict(input=x)) |
| 112 | |
| 113 | |
| 114 | def main(): |
no test coverage detected