Merge pull request open-mmlab#76 from cxxgtxy/master

Support data prefetching on the linear classification task

configs/benchmarks/linear_classification/imagenet/r50_last.py

@@ -28,15 +28,16 @@
 train_pipeline = [
     dict(type='RandomResizedCrop', size=224),
     dict(type='RandomHorizontalFlip'),
-    dict(type='ToTensor'),
-    dict(type='Normalize', **img_norm_cfg),
 ]
 test_pipeline = [
     dict(type='Resize', size=256),
     dict(type='CenterCrop', size=224),
-    dict(type='ToTensor'),
-    dict(type='Normalize', **img_norm_cfg),
 ]
+# prefetch
+prefetch = False
+if not prefetch:
+    train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)])
+    test_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)])
 data = dict(
     imgs_per_gpu=32,  # total 32*8=256, 8GPU linear cls
     workers_per_gpu=5,
@@ -45,12 +46,14 @@
         data_source=dict(
             list_file=data_train_list, root=data_train_root,
             **data_source_cfg),
-        pipeline=train_pipeline),
+        pipeline=train_pipeline,
+        prefetch=prefetch),
     val=dict(
         type=dataset_type,
         data_source=dict(
             list_file=data_test_list, root=data_test_root, **data_source_cfg),
-        pipeline=test_pipeline))
+        pipeline=test_pipeline,
+        prefetch=prefetch))
 # additional hooks
 custom_hooks = [
     dict(
@@ -60,6 +63,8 @@
         interval=1,
         imgs_per_gpu=128,
         workers_per_gpu=4,
+        prefetch=prefetch,
+        img_norm_cfg=img_norm_cfg,
         eval_param=dict(topk=(1, 5)))
 ]
 # optimizer

configs/benchmarks/linear_classification/imagenet/r50_last_sobel.py

@@ -28,15 +28,16 @@
 train_pipeline = [
     dict(type='RandomResizedCrop', size=224),
     dict(type='RandomHorizontalFlip'),
-    dict(type='ToTensor'),
-    dict(type='Normalize', **img_norm_cfg),
 ]
 test_pipeline = [
     dict(type='Resize', size=256),
     dict(type='CenterCrop', size=224),
-    dict(type='ToTensor'),
-    dict(type='Normalize', **img_norm_cfg),
 ]
+# prefetch
+prefetch = False
+if not prefetch:
+    train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)])
+    test_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)])
 data = dict(
     imgs_per_gpu=32,  # total 32*8=256, 8GPU linear cls
     workers_per_gpu=5,
@@ -45,12 +46,14 @@
         data_source=dict(
             list_file=data_train_list, root=data_train_root,
             **data_source_cfg),
-        pipeline=train_pipeline),
+        pipeline=train_pipeline,
+        prefetch=prefetch),
     val=dict(
         type=dataset_type,
         data_source=dict(
             list_file=data_test_list, root=data_test_root, **data_source_cfg),
-        pipeline=test_pipeline))
+        pipeline=test_pipeline,
+        prefetch=prefetch))
 # additional hooks
 custom_hooks = [
     dict(
@@ -60,6 +63,8 @@
         interval=1,
         imgs_per_gpu=128,
         workers_per_gpu=4,
+        prefetch=prefetch,
+        img_norm_cfg=img_norm_cfg,
         eval_param=dict(topk=(1, 5)))
 ]
 # optimizer

configs/benchmarks/linear_classification/places205/r50_multihead.py

@@ -35,15 +35,16 @@
     dict(type='CenterCrop', size=256),
     dict(type='RandomCrop', size=224),
     dict(type='RandomHorizontalFlip'),
-    dict(type='ToTensor'),
-    dict(type='Normalize', **img_norm_cfg),
 ]
 test_pipeline = [
     dict(type='Resize', size=256),
     dict(type='CenterCrop', size=224),
-    dict(type='ToTensor'),
-    dict(type='Normalize', **img_norm_cfg),
 ]
+# prefetch
+prefetch = False
+if not prefetch:
+    train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)])
+    test_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)])
 data = dict(
     imgs_per_gpu=32,  # total 32x8=256
     workers_per_gpu=4,
@@ -52,12 +53,14 @@
         data_source=dict(
             list_file=data_train_list, root=data_train_root,
             **data_source_cfg),
-        pipeline=train_pipeline),
+        pipeline=train_pipeline,
+        prefetch=prefetch),
     val=dict(
         type=dataset_type,
         data_source=dict(
             list_file=data_test_list, root=data_test_root, **data_source_cfg),
-        pipeline=test_pipeline))
+        pipeline=test_pipeline,
+        prefetch=prefetch))
 # additional hooks
 custom_hooks = [
     dict(
@@ -67,6 +70,8 @@
         interval=10,
         imgs_per_gpu=32,
         workers_per_gpu=4,
+        prefetch=prefetch,
+        img_norm_cfg=img_norm_cfg,
         eval_param=dict(topk=(1, )))
 ]
 # optimizer

configs/benchmarks/linear_classification/places205/r50_multihead_sobel.py

@@ -35,15 +35,16 @@
     dict(type='CenterCrop', size=256),
     dict(type='RandomCrop', size=224),
     dict(type='RandomHorizontalFlip'),
-    dict(type='ToTensor'),
-    dict(type='Normalize', **img_norm_cfg),
 ]
 test_pipeline = [
     dict(type='Resize', size=256),
     dict(type='CenterCrop', size=224),
-    dict(type='ToTensor'),
-    dict(type='Normalize', **img_norm_cfg),
 ]
+# prefetch
+prefetch = False
+if not prefetch:
+    train_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)])
+    test_pipeline.extend([dict(type='ToTensor'), dict(type='Normalize', **img_norm_cfg)])
 data = dict(
     imgs_per_gpu=32,  # total 32x8=256
     workers_per_gpu=4,
@@ -52,12 +53,14 @@
         data_source=dict(
             list_file=data_train_list, root=data_train_root,
             **data_source_cfg),
-        pipeline=train_pipeline),
+        pipeline=train_pipeline,
+        prefetch=prefetch),
     val=dict(
         type=dataset_type,
         data_source=dict(
             list_file=data_test_list, root=data_test_root, **data_source_cfg),
-        pipeline=test_pipeline))
+        pipeline=test_pipeline,
+        prefetch=prefetch))
 # additional hooks
 custom_hooks = [
     dict(
@@ -67,6 +70,8 @@
         interval=10,
         imgs_per_gpu=32,
         workers_per_gpu=4,
+        prefetch=prefetch,
+        img_norm_cfg=img_norm_cfg,
         eval_param=dict(topk=(1, )))
 ]
 # optimizer

docs/GETTING_STARTED.md

@@ -92,7 +92,7 @@ prefetch = True
 3 . Replacing  Pillow with Pillow-SIMD (https://github.com/uploadcare/pillow-simd.git) to make use of SIMD command sets with modern CPU.
  ```shell
 pip uninstall pillow
-pip install Pillow-SIMD
+pip install Pillow-SIMD or CC="cc -mavx2" pip install -U --force-reinstall pillow-simd if AVX2 is available.
 ```
 We test it using MoCoV2 using a total batch size of 256 on Tesla V100. The training time per step is decreased to 0.17s from 0.23s.
 ## Benchmarks

openselfsup/datasets/classification.py

@@ -4,19 +4,22 @@
 
 from .registry import DATASETS
 from .base import BaseDataset
+from .utils import to_numpy
 
 
 @DATASETS.register_module
 class ClassificationDataset(BaseDataset):
     """Dataset for classification.
     """
 
-    def __init__(self, data_source, pipeline):
-        super(ClassificationDataset, self).__init__(data_source, pipeline)
+    def __init__(self, data_source, pipeline, prefetch=False):
+        super(ClassificationDataset, self).__init__(data_source, pipeline, prefetch)
 
     def __getitem__(self, idx):
         img, target = self.data_source.get_sample(idx)
         img = self.pipeline(img)
+        if self.prefetch:
+            img = torch.from_numpy(to_numpy(img))
         return dict(img=img, gt_label=target)
 
     def evaluate(self, scores, keyword, logger=None, topk=(1, 5)):

openselfsup/hooks/validate_hook.py

@@ -42,7 +42,10 @@ def __init__(self,
             eval_kwargs['imgs_per_gpu'],
             eval_kwargs['workers_per_gpu'],
             dist=dist_mode,
-            shuffle=False)
+            shuffle=False,
+            prefetch=eval_kwargs.get('prefetch', False),
+            img_norm_cfg=eval_kwargs.get('img_norm_cfg', dict()),
+        )
         self.dist_mode = dist_mode
         self.initial = initial
         self.interval = interval