From f7f9ab4b962d3144f65b8b0dd106004e8e01494d Mon Sep 17 00:00:00 2001 From: nihui Date: Thu, 7 Jul 2022 19:21:58 +0800 Subject: [PATCH 1/2] Update test.py --- test.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/test.py b/test.py index 4f5671f..28b4170 100644 --- a/test.py +++ b/test.py @@ -15,6 +15,7 @@ if __name__ == '__main__': parser.add_argument('--img', type=str, default='', help='The path of test image') parser.add_argument('--thresh', type=float, default=0.65, help='The path of test image') parser.add_argument('--onnx', action="store_true", default=False, help='Export onnx file') + parser.add_argument('--torchscript', action="store_true", default=False, help='Export torchscript file') parser.add_argument('--cpu', action="store_true", default=False, help='Run on cpu') opt = parser.parse_args() @@ -61,6 +62,15 @@ if __name__ == '__main__': opset_version=11, # the ONNX version to export the model to do_constant_folding=True) # whether to execute constant folding for optimization + # 导出torchscript模型 + if opt.torchscript: + import copy + model_cpu = copy.deepcopy(model).cpu() + x = torch.rand(1, 3, cfg.input_height, cfg.input_width) + mod = torch.jit.trace(model_cpu, x) + mod.save("./FastestDet.pt") + print("to convert torchscript to pnnx/ncnn: ./pnnx FastestDet.pt inputshape=[1,3,%d,%d]" % (cfg.input_height, cfg.input_height)) + # 模型推理 start = time.perf_counter() preds = model(img) From 89e38dc645324e7fe30062f5e65283e9ca3dcdc3 Mon Sep 17 00:00:00 2001 From: nihui Date: Thu, 7 Jul 2022 19:23:03 +0800 Subject: [PATCH 2/2] Update README.md --- README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/README.md b/README.md index 483c7c1..cddedaa 100644 --- a/README.md +++ b/README.md @@ -181,6 +181,11 @@ TRAIN: ``` python3 test.py --yaml configs/coco.yaml --weight weights/coco_ap05_0.250_280epoch.pth --img data/3.jpg --onnx ``` +## Export torchscript +* You can export .pt by adding the --torchscript option when executing test.py + ``` + python3 test.py --yaml configs/coco.yaml --weight weights/coco_ap05_0.250_280epoch.pth --img data/3.jpg --torchscript + ``` ## onnx-runtime * You can learn about the pre and post-processing methods of FastestDet in this Sample ```