New DetectMultiBackend() class (#5549)

* New `DetectMultiBackend()` class * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * pb to pt fix * Cleanup * explicit apply_classifier path * Cleanup2 * Cleanup3 * Cleanup4 * Cleanup5 * Cleanup6 * val.py MultiBackend inference * warmup fix * to device fix * pt fix * device fix * Val cleanup * COCO128 URL to assets * half fix * detect fix * detect fix 2 * remove half from DetectMultiBackend * training half handling * training half handling 2 * training half handling 3 * Cleanup * Fix CI error * Add torchscript _extra_files * Add TorchScript * Add CoreML * CoreML cleanup * New `DetectMultiBackend()` class * pb to pt fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup * explicit apply_classifier path * Cleanup2 * Cleanup3 * Cleanup4 * Cleanup5 * Cleanup6 * val.py MultiBackend inference * warmup fix * to device fix * pt fix * device fix * Val cleanup * COCO128 URL to assets * half fix * detect fix * detect fix 2 * remove half from DetectMultiBackend * training half handling * training half handling 2 * training half handling 3 * Cleanup * Fix CI error * Add torchscript _extra_files * Add TorchScript * Add CoreML * CoreML cleanup * revert default to pt * Add Usage examples * Cleanup val Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2021-11-09 16:45:02 +01:00
1 parent 79bca2bf64
commit 3883261143
@@ -21,6 +21,7 @@ TensorFlow.js:
 """

 import argparse
+import json
 import os
 import subprocess
 import sys
@@ -54,7 +55,9 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:'
        f = file.with_suffix('.torchscript.pt')

        ts = torch.jit.trace(model, im, strict=False)
-        (optimize_for_mobile(ts) if optimize else ts).save(f)
+        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(f, _extra_files=extra_files)

        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
    except Exception as e: