feat(tools): support torchscript export (#668)

tools/export_torchscript.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import argparse
+import os
+from loguru import logger
+
+import torch
+from torch import nn
+
+from yolox.exp import get_exp
+
+
+def make_parser():
+    parser = argparse.ArgumentParser("YOLOX torchscript deploy")
+    parser.add_argument(
+        "--output-name", type=str, default="yolox.torchscript.pt", help="output name of models"
+    )
+    parser.add_argument("--batch-size", type=int, default=1, help="batch size")
+    parser.add_argument(
+        "-f",
+        "--exp_file",
+        default=None,
+        type=str,
+        help="expriment description file",
+    )
+    parser.add_argument("-expn", "--experiment-name", type=str, default=None)
+    parser.add_argument("-n", "--name", type=str, default=None, help="model name")
+    parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path")
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+
+    return parser
+
+
+@logger.catch
+def main():
+    args = make_parser().parse_args()
+    logger.info("args value: {}".format(args))
+    exp = get_exp(args.exp_file, args.name)
+    exp.merge(args.opts)
+
+    if not args.experiment_name:
+        args.experiment_name = exp.exp_name
+
+    model = exp.get_model()
+    if args.ckpt is None:
+        file_name = os.path.join(exp.output_dir, args.experiment_name)
+        ckpt_file = os.path.join(file_name, "best_ckpt.pth")
+    else:
+        ckpt_file = args.ckpt
+
+    # load the model state dict
+    ckpt = torch.load(ckpt_file, map_location="cpu")
+
+    model.eval()
+    if "model" in ckpt:
+        ckpt = ckpt["model"]
+    model.load_state_dict(ckpt)
+    model.head.decode_in_inference = False
+
+    logger.info("loading checkpoint done.")
+    dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])
+
+    mod = torch.jit.trace(model, dummy_input)
+    mod.save(args.output_name)
+    logger.info("generated torchscript model named {}".format(args.output_name))
+
+
+if __name__ == "__main__":
+    main()