(args, cfg)
| 56 | |
| 57 | |
| 58 | def torch2onnx(args, cfg): |
| 59 | model = get_model(args, cfg).cuda() |
| 60 | |
| 61 | # speed_test(model) |
| 62 | |
| 63 | onnx_name = f'{args.model_name}.onnx' |
| 64 | torch.onnx.export(model, |
| 65 | torch.rand(1, 3, args.size, args.size).cuda(), |
| 66 | onnx_name, |
| 67 | input_names=['input'], |
| 68 | output_names=['output']) |
| 69 | |
| 70 | return model |
| 71 | |
| 72 | |
| 73 | def onnx2trt(args): |