Add custom training sample code and dataset (#256)

* add train custom dataset code
* update custom docs
* fix(evaluator): fix n_samples=0 bug for the len of dataloader is 1

docs/train_custom_data.md

@@ -35,6 +35,8 @@ 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.
 
+✧✧✧ You can download the mini-coco128 dataset by the [link](https://drive.google.com/file/d/16N3u36ycNd70m23IM7vMuRQXejAJY9Fs/view?usp=sharing), and then unzip it to the `datasets` directory. The dataset has been converted from YOLO format to COCO format, and can be used directly as a dataset for testing whether the train environment can be runned successfully.
+
 ## 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.
 
@@ -59,6 +61,8 @@ Besides, you should also overwrite the `dataset` and `evaluator`, prepared befor
 
 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.
 
+✧✧✧ You can also see the `exps/example/custom` directory for more details.
+
 ## 3. Train
 Except special cases, we always recommend to use our [COCO pretrained weights](../README.md) for initializing the model.
 
@@ -72,6 +76,12 @@ or take the `YOLOX-S` VOC training for example:
 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
 ```
 
+✧✧✧ For example:
+- If you download the [mini-coco128](https://drive.google.com/file/d/16N3u36ycNd70m23IM7vMuRQXejAJY9Fs/view?usp=sharing) and unzip it to the `datasets`, you can direct run the following training code.
+    ```bash
+    python tools/train.py -f exps/example/custom/yolox_s.py -d 8 -b 64 --fp16 -o -c /path/to/yolox_s.pth.tar
+    ```
+
 (Don't worry for the different shape of detection head between the pretrained weights and your own model, we will handle it)
 
 ## 4. Tips for Best Training Results

exps/example/custom/nano.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+import torch.nn as nn
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 0.33
+        self.width = 0.25
+        self.scale = (0.5, 1.5)
+        self.random_size = (10, 20)
+        self.test_size = (416, 416)
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+        self.enable_mixup = False
+
+        # Define yourself dataset path
+        self.data_dir = "datasets/coco128"
+        self.train_ann = "instances_train2017.json"
+        self.val_ann = "instances_val2017.json"
+
+        self.num_classes = 71
+
+    def get_model(self, sublinear=False):
+
+        def init_yolo(M):
+            for m in M.modules():
+                if isinstance(m, nn.BatchNorm2d):
+                    m.eps = 1e-3
+                    m.momentum = 0.03
+        if "model" not in self.__dict__:
+            from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
+            in_channels = [256, 512, 1024]
+            # NANO model use depthwise = True, which is main difference.
+            backbone = YOLOPAFPN(self.depth, self.width, in_channels=in_channels, depthwise=True)
+            head = YOLOXHead(self.num_classes, self.width, in_channels=in_channels, depthwise=True)
+            self.model = YOLOX(backbone, head)
+
+        self.model.apply(init_yolo)
+        self.model.head.initialize_biases(1e-2)
+        return self.model

exps/example/custom/yolox_s.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+import os
+from pathlib import Path
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.depth = 0.33
+        self.width = 0.50
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+        # Define yourself dataset path
+        self.data_dir = "datasets/coco128"
+        self.train_ann = "instances_train2017.json"
+        self.val_ann = "instances_val2017.json"
+
+        self.num_classes = 71
+
+        self.max_epoch = 300
+        self.data_num_workers = 4
+        self.eval_interval = 1

yolox/evaluators/coco_evaluator.py

@@ -83,7 +83,7 @@ class COCOEvaluator:
 
         inference_time = 0
         nms_time = 0
-        n_samples = len(self.dataloader) - 1
+        n_samples = max(len(self.dataloader) - 1, 1)
 
         if trt_file is not None:
             from torch2trt import TRTModule

yolox/exp/yolox_base.py

@@ -2,17 +2,18 @@
 # -*- coding:utf-8 -*-
 # Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 
+import os
+import random
+
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 
-import os
-import random
-
 from .base_exp import BaseExp
 
 
 class Exp(BaseExp):
+
     def __init__(self):
         super().__init__()
 
@@ -26,6 +27,7 @@ class Exp(BaseExp):
         self.data_num_workers = 4
         self.input_size = (640, 640)
         self.random_size = (14, 26)
+        self.data_dir = None
         self.train_ann = "instances_train2017.json"
         self.val_ann = "instances_val2017.json"
 
@@ -60,7 +62,7 @@ class Exp(BaseExp):
         self.nmsthre = 0.65
 
     def get_model(self):
-        from yolox.models import YOLOPAFPN, YOLOX, YOLOXHead
+        from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
 
         def init_yolo(M):
             for m in M.modules():
@@ -81,15 +83,15 @@ class Exp(BaseExp):
     def get_data_loader(self, batch_size, is_distributed, no_aug=False):
         from yolox.data import (
             COCODataset,
+            TrainTransform,
+            YoloBatchSampler,
             DataLoader,
             InfiniteSampler,
             MosaicDetection,
-            TrainTransform,
-            YoloBatchSampler
         )
 
         dataset = COCODataset(
-            data_dir=None,
+            data_dir=self.data_dir,
             json_file=self.train_ann,
             img_size=self.input_size,
             preproc=TrainTransform(
@@ -121,7 +123,9 @@ class Exp(BaseExp):
         if is_distributed:
             batch_size = batch_size // dist.get_world_size()
 
-        sampler = InfiniteSampler(len(self.dataset), seed=self.seed if self.seed else 0)
+        sampler = InfiniteSampler(
+            len(self.dataset), seed=self.seed if self.seed else 0
+        )
 
         batch_sampler = YoloBatchSampler(
             sampler=sampler,
@@ -141,7 +145,7 @@ class Exp(BaseExp):
         tensor = torch.LongTensor(2).cuda()
 
         if rank == 0:
-            size_factor = self.input_size[1] * 1.0 / self.input_size[0]
+            size_factor = self.input_size[1] * 1. / self.input_size[0]
             size = random.randint(*self.random_size)
             size = (int(32 * size), 32 * int(size * size_factor))
             tensor[0] = size[0]
@@ -186,7 +190,6 @@ class Exp(BaseExp):
 
     def get_lr_scheduler(self, lr, iters_per_epoch):
         from yolox.utils import LRScheduler
-
         scheduler = LRScheduler(
             self.scheduler,
             lr,
@@ -203,7 +206,7 @@ class Exp(BaseExp):
         from yolox.data import COCODataset, ValTransform
 
         valdataset = COCODataset(
-            data_dir=None,
+            data_dir=self.data_dir,
             json_file=self.val_ann if not testdev else "image_info_test-dev2017.json",
             name="val2017" if not testdev else "test2017",
             img_size=self.test_size,