chore(tools): fix potential bug and update docs (#89)

.gitignore

@@ -1,6 +1,10 @@
 ### Linux ###
 *~
 
+# user experiments directory
+YOLOX_outputs/
+datasets/
+
 # temporary files which can be created if a process still has a handle open of a deleted file
 .fuse_hidden*
 

docs/train_custom_data.md

@@ -7,11 +7,11 @@ We take an example of fine-tuning YOLOX-S model on VOC dataset to give a more cl
 Clone this repo and follow the [README](../README.md) to install YOLOX.
 
 ## 1. Create your own dataset
-**Step 1** Prepare your own dataset with images and labels first. For labeling images, you may use a tool like [Labelme](https://github.com/wkentaro/labelme) or [CVAT](https://github.com/openvinotoolkit/cvat).
+**Step 1** Prepare your own dataset with images and labels first. For labeling images, you can use tools like [Labelme](https://github.com/wkentaro/labelme) or [CVAT](https://github.com/openvinotoolkit/cvat).
 
-**Step 2** Then, you should write the corresponding Dataset Class which can load images and labels through "\_\_getitem\_\_" method. We currently support COCO format and VOC format.
+**Step 2** Then, you should write the corresponding Dataset Class which can load images and labels through `__getitem__` method. We currently support COCO format and VOC format.
 
-You can also write the Dataset by you own. Let's take the [VOC](../yolox/data/datasets/voc.py#L151) Dataset file for example:
+You can also write the Dataset by your own. Let's take the [VOC](../yolox/data/datasets/voc.py#L151) Dataset file for example:
 ```python
     @Dataset.resize_getitem
     def __getitem__(self, index):
@@ -23,27 +23,28 @@ You can also write the Dataset by you own. Let's take the [VOC](../yolox/data/da
         return img, target, img_info, img_id
 ```
 
-One more thing worth noting is that you should also implement "[pull_item](../yolox/data/datasets/voc.py#L129)" and "[load_anno](../yolox/data/datasets/voc.py#L121)" method for the Mosiac and MixUp augmentation.
+One more thing worth noting is that you should also implement [pull_item](../yolox/data/datasets/voc.py#L129) and [load_anno](../yolox/data/datasets/voc.py#L121) method for the `Mosiac` and `MixUp` augmentations.
 
 **Step 3** Prepare the evaluator. We currently have [COCO evaluator](../yolox/evaluators/coco_evaluator.py) and [VOC evaluator](../yolox/evaluators/voc_evaluator.py).
-If you have your own format data or evaluation metric, you may write your own evaluator.
+If you have your own format data or evaluation metric, you can write your own evaluator.
+
+**Step 4** Put your dataset under `$YOLOX_DIR/datasets`, for VOC:
 
-**Step 4** Put your dataset under $YOLOX_DIR/datasets$, for VOC:
 ```shell
 ln -s /path/to/your/VOCdevkit ./datasets/VOCdevkit
 ```
-* The path "VOCdevkit" will be used in your exp file described in next section.Specifically, in "get_data_loader" and "get_eval_loader" function.
+* The path "VOCdevkit" will be used in your exp file described in next section. Specifically, in `get_data_loader` and `get_eval_loader` function.
 
 ## 2. Create your Exp file to control everything
 We put everything involved in a model to one single Exp file, including model setting, training setting, and testing setting.
 
 **A complete Exp file is at [yolox_base.py](../yolox/exp/yolox_base.py).** It may be too long to write for every exp, but you can inherit the base Exp file and only overwrite the changed part.
 
-Let's still take the [VOC Exp file](../exps/example/yolox_voc/yolox_voc_s.py) for an example.
+Let's take the [VOC Exp file](../exps/example/yolox_voc/yolox_voc_s.py) as an example.
 
-We select YOLOX-S model here, so we should change the network depth and width. VOC has only 20 classes, so we should also change the num_classes.
+We select `YOLOX-S` model here, so we should change the network depth and width. VOC has only 20 classes, so we should also change the `num_classes`.
 
-These configs are changed in the init() methd:
+These configs are changed in the `init()` method:
 ```python
 class Exp(MyExp):
     def __init__(self):
@@ -54,19 +55,19 @@ class Exp(MyExp):
         self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
 ```
 
-Besides, you should also overwrite the dataset and evaluator preprared before to training the model on your own data.
+Besides, you should also overwrite the `dataset` and `evaluator`, prepared before training the model on your own data.
 
-Please see "[get_data_loader](../exps/example/yolox_voc/yolox_voc_s.py#L20)", "[get_eval_loader](../exps/example/yolox_voc/yolox_voc_s.py#L82)", and "[get_evaluator](../exps/example/yolox_voc/yolox_voc_s.py#L113)" for more details.
+Please see [get_data_loader](../exps/example/yolox_voc/yolox_voc_s.py#L20), [get_eval_loader](../exps/example/yolox_voc/yolox_voc_s.py#L82), and [get_evaluator](../exps/example/yolox_voc/yolox_voc_s.py#L113) for more details.
 
 ## 3. Train
-Except special cases, we always recommend to use our [COCO pretrained weights](../README.md) for initializing.
+Except special cases, we always recommend to use our [COCO pretrained weights](../README.md) for initializing the model.
 
-Once you get the Exp file and the COCO pretrained weights we provided, you can train your own model by the following command:
+Once you get the Exp file and the COCO pretrained weights we provided, you can train your own model by the following below command:
 ```bash
 python tools/train.py -f /path/to/your/Exp/file -d 8 -b 64 --fp16 -o -c /path/to/the/pretrained/weights
 ```
 
-or take the YOLOX-S VOC training for example:
+or take the `YOLOX-S` VOC training for example:
 ```bash
 python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 8 -b 64 --fp16 -o -c /path/to/yolox_s.pth.tar
 ```
@@ -75,16 +76,17 @@ python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 8 -b 64 --fp16
 
 ## 4. Tips for Best Training Results
 
-As YOLOX is an anchor-free detector with only several hyper-parameters, most of the time good results can be obtained with no changes to the models or training settings.
+As **YOLOX** is an anchor-free detector with only several hyper-parameters, most of the time good results can be obtained with no changes to the models or training settings.
 We thus always recommend you first train with all default training settings.
 
-If at first you don't get good results, there are steps you could consider to take to improve.
+If at first you don't get good results, there are steps you could consider to improve the model.
 
-**Model Selection** We provide YOLOX-Nano, YOLOX-Tiny, and YOLOX-S for mobile deployments, while YOLOX-M/L/X for cloud or high performance GPU deployments.
+**Model Selection** We provide `YOLOX-Nano`, `YOLOX-Tiny`, and `YOLOX-S` for mobile deployments, while `YOLOX-M`/`L`/`X` for cloud or high performance GPU deployments.
 
-If your deployment meets some trouble of compatibility. we recommand YOLOX-DarkNet53.
+If your deployment meets any compatibility issues. we recommend `YOLOX-DarkNet53`.
 
 **Training Configs** If your training overfits early, then you can reduce max\_epochs or decrease the base\_lr and min\_lr\_ratio in your Exp file:
+
 ```python
 # --------------  training config --------------------- #
     self.warmup_epochs = 5

requirements.txt

@@ -1,3 +1,4 @@
+# TODO: Update with exact module version
 numpy
 torch>=1.7
 opencv_python
@@ -10,4 +11,8 @@ thop
 ninja
 tabulate
 tensorboard
-onnxruntime
+
+# verified versions
+onnx==1.8.1
+onnxruntime==1.8.0
+onnx-simplifier==0.3.5

tools/export_onnx.py

@@ -16,30 +16,25 @@ from yolox.utils import replace_module
 
 def make_parser():
     parser = argparse.ArgumentParser("YOLOX onnx deploy")
-    parser.add_argument(
-        "--output-name", type=str, default="yolox.onnx", help="output name of models"
-    )
-    parser.add_argument("--input", default="images", type=str, help="input name of onnx model")
-    parser.add_argument("--output", default="output", type=str, help="output name of onnx model")
-    parser.add_argument("-o", "--opset", default=11, type=int, help="onnx opset version")
-    parser.add_argument("--no-onnxsim", action="store_true", help="use onnxsim or not")
-
-    parser.add_argument(
-        "-f",
-        "--exp_file",
-        default=None,
-        type=str,
-        help="expriment description file",
-    )
+    parser.add_argument("--output-name", type=str, default="yolox.onnx",
+                        help="output name of models")
+    parser.add_argument("--input", default="images", type=str,
+                        help="input node name of onnx model")
+    parser.add_argument("--output", default="output", type=str,
+                        help="output node name of onnx model")
+    parser.add_argument("-o", "--opset", default=11, type=int,
+                        help="onnx opset version")
+    parser.add_argument("--no-onnxsim", action="store_true",
+                        help="use onnxsim or not")
+    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,
-    )
+    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
 
@@ -61,8 +56,8 @@ def main():
     else:
         ckpt_file = args.ckpt
 
-    ckpt = torch.load(ckpt_file, map_location="cpu")
     # load the model state dict
+    ckpt = torch.load(ckpt_file, map_location="cpu")
 
     model.eval()
     if "model" in ckpt:
@@ -71,7 +66,7 @@ def main():
     model = replace_module(model, nn.SiLU, SiLU)
     model.head.decode_in_inference = False
 
-    logger.info("loaded checkpoint done.")
+    logger.info("loading checkpoint done.")
     dummy_input = torch.randn(1, 3, exp.test_size[0], exp.test_size[1])
     torch.onnx._export(
         model,
@@ -81,12 +76,18 @@ def main():
         output_names=[args.output],
         opset_version=args.opset,
     )
-    logger.info("generate onnx named {}".format(args.output_name))
+    logger.info("generated onnx model named {}".format(args.output_name))
 
     if not args.no_onnxsim:
+        import onnx
+        from onnxsim import simplify
+
         # use onnxsimplify to reduce reduent model.
-        os.system("python3 -m onnxsim {} {}".format(args.output_name, args.output_name))
-        logger.info("generate simplify onnx named {}".format(args.output_name))
+        onnx_model = onnx.load(args.output_name)
+        model_simp, check = simplify(onnx_model)
+        assert check, "Simplified ONNX model could not be validated"
+        onnx.save(model_simp, args.output_name)
+        logger.info("generated simplified onnx model named {}".format(args.output_name))
 
 
 if __name__ == "__main__":