feat(data): support custom dataset cache (Megvii-BaseDetection#1584)

feat(data): support custom dataset cache

Author: yuangpeng

README.md

@@ -122,9 +122,9 @@ python -m yolox.tools.train -n yolox-s -d 8 -b 64 --fp16 -o [--cache]
 * -d: number of gpu devices
 * -b: total batch size, the recommended number for -b is num-gpu * 8
 * --fp16: mixed precision training
-* --cache: caching imgs into RAM to accelarate training, which need large system RAM. 
+* --cache: caching imgs into RAM to accelarate training, which need large system RAM.
+
 
-  
 
 When using -f, the above commands are equivalent to:
 ```shell
@@ -140,7 +140,8 @@ We also support multi-nodes training. Just add the following args:
 * --num\_machines: num of your total training nodes
 * --machine\_rank: specify the rank of each node
 
-Suppose you want to train YOLOX on 2 machines, and your master machines's IP is 123.123.123.123, use port 12312 and TCP.  
+Suppose you want to train YOLOX on 2 machines, and your master machines's IP is 123.123.123.123, use port 12312 and TCP.
+
 On master machine, run
 ```shell
 python tools/train.py -n yolox-s -b 128 --dist-url tcp://123.123.123.123:12312 --num_machines 2 --machine_rank 0
@@ -163,7 +164,8 @@ python tools/train.py -n yolox-s -d 8 -b 64 --fp16 -o [--cache] --logger wandb w
 
 An example wandb dashboard is available [here](https://wandb.ai/manan-goel/yolox-nano/runs/3pzfeom0)
 
-**Others**  
+**Others**
+
 See more information with the following command:
 ```shell
 python -m yolox.tools.train --help
@@ -202,6 +204,7 @@ python -m yolox.tools.eval -n  yolox-s -c yolox_s.pth -b 1 -d 1 --conf 0.001 --f
 <summary>Tutorials</summary>
 
 *  [Training on custom data](docs/train_custom_data.md)
+*  [Caching for custom data](docs/cache.md)
 *  [Manipulating training image size](docs/manipulate_training_image_size.md)
 *  [Freezing model](docs/freeze_module.md)
 

docs/cache.md

@@ -0,0 +1,97 @@
+# Cache Custom Data
+
+The caching feature is specifically tailored for users with ample memory resources. However, we still offer the option to cache data to disk, but disk performance can vary and may not guarantee optimal user experience. Implementing custom dataset RAM caching is also more straightforward and user-friendly compared to disk caching. With a few simple modifications, users can expect to see a significant increase in training speed, with speeds nearly double that of non-cached datasets.
+
+This page explains how to cache your own custom data with YOLOX.
+
+## 0. Before you start
+
+**Step1** Clone this repo and follow the [README](../README.md) to install YOLOX.
+
+**Stpe2** Read the [Training on custom data](./train_custom_data.md) tutorial to understand how to prepare your custom data.
+
+## 1. Inheirit from `CacheDataset`
+
+
+**Step1** Create a custom dataset that inherits from the `CacheDataset` class. Note that whether inheriting from `Dataset` or `CacheDataset `, the `__init__()` method of your custom dataset should take the following keyword arguments: `input_dimension`, `cache`, and `cache_type`. Also, call `super().__init__()` and pass in `input_dimension`, `num_imgs`, `cache`, and `cache_type` as input, where `num_imgs` is the size of the dataset.
+
+**Step2** Implement the abstract function `read_img(self, index, use_cache=True)` of parent class and decorate it with `@cache_read_img`.  This function takes an `index` as input and returns an `image`, and the returned image will be used for caching. It is recommended to put all repetitive and fixed post-processing operations on the image in this function to reduce the post-processing time of the image during training.
+
+```python
+# CustomDataset.py
+from yolox.data.datasets import CacheDataset, cache_read_img
+
+class CustomDataset(CacheDataset):
+    def __init__(self, input_dimension, cache, cache_type, *args, **kwargs):
+        # Get the required keyword arguments of super().__init__()
+        super().__init__(
+            input_dimension=input_dimension,
+            num_imgs=num_imgs,
+            cache=cache,
+            cache_type=cache_type
+        )
+        # ...
+
+    @cache_read_img
+    def read_img(self, index, use_cache=True):
+        # get image ...
+        # (optional) repetitive and fixed post-processing operations for image
+        return image
+```
+
+## 2. Create your Exp file and return your custom dataset
+
+**Step1** Create a new class that inherits from the `Exp` class provided by the `yolox_base.py`. Override the `get_dataset()` and `get_eval_dataset()` method to return an instance of your custom dataset.
+
+**Step2** Implement your own `get_evaluator` method to return an instance of your custom evaluator.
+
+```python
+# CustomeExp.py
+from yolox.exp import Exp as MyExp
+
+class Exp(MyExp):
+    def get_dataset(self, cache, cache_type: str = "ram"):
+        return CustomDataset(
+            input_dimension=self.input_size,
+            cache=cache,
+            cache_type=cache_type
+        )
+
+    def get_eval_dataset(self):
+        return CustomDataset(
+            input_dimension=self.input_size,
+        )
+
+    def get_evaluator(self, batch_size, is_distributed, testdev=False, legacy=False):
+        return CustomEvaluator(
+            dataloader=self.get_eval_loader(batch_size, is_distributed, testdev=testdev, legacy=legacy),
+            img_size=self.test_size,
+            confthre=self.test_conf,
+            nmsthre=self.nmsthre,
+            num_classes=self.num_classes,
+            testdev=testdev,
+        )
+```
+
+**(Optional)** `get_data_loader` and `get_eval_loader` are now a default behavior in `yolox_base.py` and generally do not need to be changed. If you have to change `get_data_loader`, you need to add the following code at the beginning.
+
+```python
+# CustomeExp.py
+from yolox.exp import Exp as MyExp
+
+class Exp(MyExp):
+    def get_data_loader(self, batch_size, is_distributed, no_aug=False, cache_img: str = None):
+        if self.dataset is None:
+            with wait_for_the_master():
+                assert cache_img is None
+                self.dataset = self.get_dataset(cache=False, cache_type=cache_img)
+        # ...
+
+```
+
+## 3. Cache to Disk
+It's important to note that the `cache_type` can be `"ram"` or `"disk"`, depending on where you want to cache your dataset. If you choose `"disk"`, you need to pass in additional parameters to `super().__init__()` of `CustomDataset`: `data_dir`, `cache_dir_name`, `path_filename`.
+
+- `data_dir`: the root directory of the dataset, e.g. `/path/to/COCO`.
+- `cache_dir_name`: the name of the directory to cache to disk, for example `"custom_cache"`, then the files cached to disk will be saved under `/path/to/COCO/custom_cache`.
+- `path_filename`: a list of paths to the data relative to the `data_dir`, e.g. if you have data `/path/to/COCO/train/1.jpg`, `/path/to/COCO/train/2.jpg`, then `path_filename = ['train/1.jpg', ' train/2.jpg']`.

exps/example/yolox_voc/yolox_voc_s.py

@@ -1,9 +1,6 @@
 # encoding: utf-8
 import os
 
-import torch
-import torch.distributed as dist
-
 from yolox.data import get_yolox_datadir
 from yolox.exp import Exp as MyExp
 
@@ -24,115 +21,40 @@ def __init__(self):
 
         self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
 
-    def get_data_loader(self, batch_size, is_distributed, no_aug=False, cache_img=False):
-        from yolox.data import (
-            VOCDetection,
-            TrainTransform,
-            YoloBatchSampler,
-            DataLoader,
-            InfiniteSampler,
-            MosaicDetection,
-            worker_init_reset_seed,
-        )
-        from yolox.utils import (
-            wait_for_the_master,
-            get_local_rank,
-        )
-        local_rank = get_local_rank()
+    def get_dataset(self, cache: bool, cache_type: str = "ram"):
+        from yolox.data import VOCDetection, TrainTransform
 
-        with wait_for_the_master(local_rank):
-            dataset = VOCDetection(
-                data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
-                image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
-                img_size=self.input_size,
-                preproc=TrainTransform(
-                    max_labels=50,
-                    flip_prob=self.flip_prob,
-                    hsv_prob=self.hsv_prob),
-                cache=cache_img,
-            )
-
-        dataset = MosaicDetection(
-            dataset,
-            mosaic=not no_aug,
+        return VOCDetection(
+            data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
+            image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
             img_size=self.input_size,
             preproc=TrainTransform(
-                max_labels=120,
+                max_labels=50,
                 flip_prob=self.flip_prob,
                 hsv_prob=self.hsv_prob),
-            degrees=self.degrees,
-            translate=self.translate,
-            mosaic_scale=self.mosaic_scale,
-            mixup_scale=self.mixup_scale,
-            shear=self.shear,
-            enable_mixup=self.enable_mixup,
-            mosaic_prob=self.mosaic_prob,
-            mixup_prob=self.mixup_prob,
-        )
-
-        self.dataset = dataset
-
-        if is_distributed:
-            batch_size = batch_size // dist.get_world_size()
-
-        sampler = InfiniteSampler(
-            len(self.dataset), seed=self.seed if self.seed else 0
+            cache=cache,
+            cache_type=cache_type,
         )
 
-        batch_sampler = YoloBatchSampler(
-            sampler=sampler,
-            batch_size=batch_size,
-            drop_last=False,
-            mosaic=not no_aug,
-        )
-
-        dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True}
-        dataloader_kwargs["batch_sampler"] = batch_sampler
-
-        # Make sure each process has different random seed, especially for 'fork' method
-        dataloader_kwargs["worker_init_fn"] = worker_init_reset_seed
-
-        train_loader = DataLoader(self.dataset, **dataloader_kwargs)
-
-        return train_loader
-
-    def get_eval_loader(self, batch_size, is_distributed, testdev=False, legacy=False):
+    def get_eval_dataset(self, **kwargs):
         from yolox.data import VOCDetection, ValTransform
+        legacy = kwargs.get("legacy", False)
 
-        valdataset = VOCDetection(
+        return VOCDetection(
             data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
             image_sets=[('2007', 'test')],
             img_size=self.test_size,
             preproc=ValTransform(legacy=legacy),
         )
 
-        if is_distributed:
-            batch_size = batch_size // dist.get_world_size()
-            sampler = torch.utils.data.distributed.DistributedSampler(
-                valdataset, shuffle=False
-            )
-        else:
-            sampler = torch.utils.data.SequentialSampler(valdataset)
-
-        dataloader_kwargs = {
-            "num_workers": self.data_num_workers,
-            "pin_memory": True,
-            "sampler": sampler,
-        }
-        dataloader_kwargs["batch_size"] = batch_size
-        val_loader = torch.utils.data.DataLoader(valdataset, **dataloader_kwargs)
-
-        return val_loader
-
     def get_evaluator(self, batch_size, is_distributed, testdev=False, legacy=False):
         from yolox.evaluators import VOCEvaluator
 
-        val_loader = self.get_eval_loader(batch_size, is_distributed, testdev, legacy)
-        evaluator = VOCEvaluator(
-            dataloader=val_loader,
+        return VOCEvaluator(
+            dataloader=self.get_eval_loader(batch_size, is_distributed,
+                                            testdev=testdev, legacy=legacy),
             img_size=self.test_size,
             confthre=self.test_conf,
             nmsthre=self.nmsthre,
             num_classes=self.num_classes,
         )
-        return evaluator

tools/train.py

@@ -131,7 +131,7 @@ def main(exp: Exp, args):
     assert num_gpu <= get_num_devices()
 
     if args.cache is not None:
-        exp.create_cache_dataset(args.cache)
+        exp.dataset = exp.get_dataset(cache=True, cache_type=args.cache)
 
     dist_url = "auto" if args.dist_url is None else args.dist_url
     launch(

yolox/data/datasets/__init__.py

@@ -4,6 +4,6 @@
 
 from .coco import COCODataset
 from .coco_classes import COCO_CLASSES
-from .datasets_wrapper import ConcatDataset, Dataset, MixConcatDataset
+from .datasets_wrapper import CacheDataset, ConcatDataset, Dataset, MixConcatDataset
 from .mosaicdetection import MosaicDetection
 from .voc import VOCDetection