| |
| |
| |
|
|
| import argparse |
| import os |
| import shutil |
| from loguru import logger |
|
|
| import tensorrt as trt |
| import torch |
| from torch2trt import torch2trt |
|
|
| from yolox.exp import get_exp |
|
|
|
|
| def make_parser(): |
| parser = argparse.ArgumentParser("YOLOX ncnn deploy") |
| 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( |
| "-f", |
| "--exp_file", |
| default=None, |
| type=str, |
| help="please input your experiment description file", |
| ) |
| parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path") |
| parser.add_argument( |
| "-w", '--workspace', type=int, default=32, help='max workspace size in detect' |
| ) |
| parser.add_argument("-b", '--batch', type=int, default=1, help='max batch size in detect') |
| return parser |
|
|
|
|
| @logger.catch |
| @torch.no_grad() |
| def main(): |
| args = make_parser().parse_args() |
| exp = get_exp(args.exp_file, args.name) |
| if not args.experiment_name: |
| args.experiment_name = exp.exp_name |
|
|
| model = exp.get_model() |
| file_name = os.path.join(exp.output_dir, args.experiment_name) |
| os.makedirs(file_name, exist_ok=True) |
| if args.ckpt is None: |
| ckpt_file = os.path.join(file_name, "best_ckpt.pth") |
| else: |
| ckpt_file = args.ckpt |
|
|
| ckpt = torch.load(ckpt_file, map_location="cpu") |
| |
|
|
| model.load_state_dict(ckpt["model"]) |
| logger.info("loaded checkpoint done.") |
| model.eval() |
| model.cuda() |
| model.head.decode_in_inference = False |
| x = torch.ones(1, 3, exp.test_size[0], exp.test_size[1]).cuda() |
| model_trt = torch2trt( |
| model, |
| [x], |
| fp16_mode=True, |
| log_level=trt.Logger.INFO, |
| max_workspace_size=(1 << args.workspace), |
| max_batch_size=args.batch, |
| ) |
| torch.save(model_trt.state_dict(), os.path.join(file_name, "model_trt.pth")) |
| logger.info("Converted TensorRT model done.") |
| engine_file = os.path.join(file_name, "model_trt.engine") |
| engine_file_demo = os.path.join("demo", "TensorRT", "cpp", "model_trt.engine") |
| with open(engine_file, "wb") as f: |
| f.write(model_trt.engine.serialize()) |
|
|
| shutil.copyfile(engine_file, engine_file_demo) |
|
|
| logger.info("Converted TensorRT model engine file is saved for C++ inference.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
