From 024d2418ec127bbc7e93f1bbdfbf2b72713a9917 Mon Sep 17 00:00:00 2001
From: Emmanuel Benazera <emmanuel.benazera@jolibrain.com>
Date: Thu, 27 Jan 2022 18:15:46 +0100
Subject: [PATCH] chore: yolox export filenames and documentation

---
 tools/torch/README.md      | 20 ++++++++++++++++++++
 tools/torch/trace_yolox.py |  2 +-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/tools/torch/README.md b/tools/torch/README.md
index ce2b70dced..6c7767e8f1 100644
--- a/tools/torch/README.md
+++ b/tools/torch/README.md
@@ -16,3 +16,23 @@ Utility script to trace NLP models from Huggingface's pytorch-transformers. Requ
 pip3 install --user pytorch-transformers==1.1
 ```
 At the moment a CUDA model can not be convertible to CPU and vice versa. This may change with a future version of pytorch-transformers.
+
+* `trace_yolox.py`
+
+Requirements:
+
+```
+pip3 install thop
+pip3 install loguru
+```
+and custom YOLOX repository:
+```
+git clone https://github.com/jolibrain/YOLOX.git
+git checkout jit_script
+```
+
+Then to export a pretrained model before using it with DeepDetect:
+
+```
+python3 trace_yolox.py yolox-m -o /path/to/output/ --yolox_path /path/to/YOLOX/ --backbone_weights /path/to/pretrained/yolox_m.pth --num_classes 2 --img_width 300 --img_height 300
+```
\ No newline at end of file
diff --git a/tools/torch/trace_yolox.py b/tools/torch/trace_yolox.py
index a75dfa7b58..81a57ac588 100755
--- a/tools/torch/trace_yolox.py
+++ b/tools/torch/trace_yolox.py
@@ -104,7 +104,7 @@ def main():
         model.to(device)
         model.eval()
 
-        filename += ".pt"
+        filename += "-" + str(args.img_width) + 'x' + str(args.img_height) + "_cls" + str(args.num_classes) + ".pt"
         script_module = torch.jit.script(model)
         logging.info("Save jit model at %s" % filename)
         script_module.save(filename)