From 5fb246de08b1ac1479495ae7f0bfb5e5d286f799 Mon Sep 17 00:00:00 2001 From: yinrong Date: Tue, 4 Jan 2022 11:50:39 +0800 Subject: [PATCH 1/2] fix export of TorchScript on mobile --- export.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/export.py b/export.py index 3b677d2ca144..45b14e68376c 100644 --- a/export.py +++ b/export.py @@ -75,7 +75,13 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:' ts = torch.jit.trace(model, im, strict=False) d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names} extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap() - (optimize_for_mobile(ts) if optimize else ts).save(str(f), _extra_files=extra_files) + if optimize: + ts = optimize_for_mobile(ts) + # pytorch BETA + # https://pytorch.org/tutorials/recipes/mobile_interpreter.html + ts._save_for_lite_interpreter(str(f), _extra_files=extra_files) + else: + ts.save(str(f), _extra_files=extra_files) LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') except Exception as e: From a35aa72bd375341170bdc0d614bfa848202bdb98 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 3 Jan 2022 20:22:09 -0800 Subject: [PATCH 2/2] Cleanup --- export.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/export.py b/export.py index 45b14e68376c..d169aa0de5f5 100644 --- a/export.py +++ b/export.py @@ -75,11 +75,8 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:' ts = torch.jit.trace(model, im, strict=False) d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names} extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap() - if optimize: - ts = optimize_for_mobile(ts) - # pytorch BETA - # https://pytorch.org/tutorials/recipes/mobile_interpreter.html - ts._save_for_lite_interpreter(str(f), _extra_files=extra_files) + if optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html + optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files) else: ts.save(str(f), _extra_files=extra_files)