Add mAP Metrics and validation (#89)

* Refactor YOLOModule.training_step

* Introduce coco evaluator metrics

* Refactor dataloader getter in UnitTest

* Rename CocoDetection to COCODetection

* Refactor vanilla dataset getter

* Bump PL to 1.3.0rc1

* Fix COCOEvaluator metrics

* Remove unused codes

* Assert mAP on coco128 greater than 0.41

* Refactor COCOEvaluator

* Rename _utils.py to distributed.py

* Fix COCO evaluation metrics with Lightning trainer

.gitignore

@@ -102,6 +102,5 @@ doxygen
 *.sst
 *.ldb
 LOCK
-LOG*
 CURRENT
 MANIFEST-*

requirements.txt

@@ -11,4 +11,6 @@ onnx>=1.8.0
 ipython
 # pycocotools need python3.7 as minimal
 # pip install -U pycocotools>=2.0.2  # corresponds to https://github.com/ppwwyyxx/cocoapi
-pytorch-lightning>=1.1.1
+pytorch_lightning==1.3.0rc1
+torchmetrics
+tabulate

test/test_data_pipeline.py

@@ -2,42 +2,26 @@
 from pathlib import Path
 import unittest
 
-import torch.utils.data
 from torch import Tensor
 
 from yolort.data import DetectionDataModule
-from yolort.data.coco import CocoDetection
-from yolort.data.transforms import collate_fn, default_train_transforms
-from yolort.utils import prepare_coco128
-
-from .dataset_utils import DummyCOCODetectionDataset
+import yolort.data._helper as data_helper
 
 from typing import Dict
 
 
 class DataPipelineTester(unittest.TestCase):
-    def test_vanilla_dataloader(self):
+    def test_vanilla_dataset(self):
         # Acquire the images and labels from the coco128 dataset
-        data_path = Path('data-bin')
-        coco128_dirname = 'coco128'
-        coco128_path = data_path / coco128_dirname
-        image_root = coco128_path / 'images' / 'train2017'
-        annotation_file = coco128_path / 'annotations' / 'instances_train2017.json'
-
-        if not annotation_file.is_file():
-            prepare_coco128(data_path, dirname=coco128_dirname)
-
-        dataset = CocoDetection(image_root, annotation_file, default_train_transforms())
+        dataset = data_helper.get_dataset(data_root='data-bin', mode='train')
         # Test the datasets
         image, target = next(iter(dataset))
         self.assertIsInstance(image, Tensor)
         self.assertIsInstance(target, Dict)
 
-        batch_size = 4
-        sampler = torch.utils.data.RandomSampler(dataset)
-        batch_sampler = torch.utils.data.BatchSampler(sampler, batch_size, drop_last=True)
-        data_loader = torch.utils.data.DataLoader(
-            dataset, batch_sampler=batch_sampler, collate_fn=collate_fn, num_workers=0)
+    def test_vanilla_dataloader(self):
+        batch_size = 8
+        data_loader = data_helper.get_dataloader(data_root='data-bin', mode='train', batch_size=batch_size)
         # Test the dataloader
         images, targets = next(iter(data_loader))
 
@@ -54,7 +38,7 @@ def test_vanilla_dataloader(self):
     def test_detection_data_module(self):
         # Setup the DataModule
         batch_size = 4
-        train_dataset = DummyCOCODetectionDataset(num_samples=128)
+        train_dataset = data_helper.DummyCOCODetectionDataset(num_samples=128)
         data_module = DetectionDataModule(train_dataset, batch_size=batch_size)
         self.assertEqual(data_module.batch_size, batch_size)
 
@@ -72,6 +56,6 @@ def test_detection_data_module(self):
     def test_prepare_coco128(self):
         data_path = Path('data-bin')
         coco128_dirname = 'coco128'
-        prepare_coco128(data_path, dirname=coco128_dirname)
+        data_helper.prepare_coco128(data_path, dirname=coco128_dirname)
         annotation_file = data_path / coco128_dirname / 'annotations' / 'instances_train2017.json'
         self.assertTrue(annotation_file.is_file())

test/test_engine.py

@@ -1,6 +1,6 @@
 # Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
-from pathlib import Path
 import unittest
+from pathlib import Path
 
 import torch
 from torch import Tensor
@@ -9,17 +9,12 @@
 
 import pytorch_lightning as pl
 
-from yolort.data import DetectionDataModule
-from yolort.data.coco import CocoDetection
-from yolort.data.transforms import collate_fn, default_train_transforms
+from yolort.data import COCOEvaluator, DetectionDataModule
+import yolort.data._helper as data_helper
 
+from yolort.models import yolov5s
 from yolort.models.yolo import yolov5_darknet_pan_s_r31
 from yolort.models.transform import nested_tensor_from_tensor_list
-from yolort.models import yolov5s
-
-from yolort.utils import prepare_coco128
-
-from .dataset_utils import DummyCOCODetectionDataset
 
 from typing import Dict
 
@@ -65,26 +60,9 @@ def test_train_with_vanilla_module(self):
         # Define the device
         device = torch.device('cpu')
 
-        # Prepare the datasets for training
-        # Acquire the images and labels from the coco128 dataset
-        data_path = Path('data-bin')
-        coco128_dirname = 'coco128'
-        coco128_path = data_path / coco128_dirname
-        image_root = coco128_path / 'images' / 'train2017'
-        annotation_file = coco128_path / 'annotations' / 'instances_train2017.json'
-
-        if not annotation_file.is_file():
-            prepare_coco128(data_path, dirname=coco128_dirname)
-
-        batch_size = 4
-
-        dataset = CocoDetection(image_root, annotation_file, default_train_transforms())
-        sampler = torch.utils.data.RandomSampler(dataset)
-        batch_sampler = torch.utils.data.BatchSampler(sampler, batch_size, drop_last=True)
-        data_loader = torch.utils.data.DataLoader(
-            dataset, batch_sampler=batch_sampler, collate_fn=collate_fn, num_workers=0)
+        train_dataloader = data_helper.get_dataloader(data_root='data-bin', mode='train')
         # Sample a pair of images/targets
-        images, targets = next(iter(data_loader))
+        images, targets = next(iter(train_dataloader))
         images = [img.to(device) for img in images]
         targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
 
@@ -98,9 +76,9 @@ def test_train_with_vanilla_module(self):
         self.assertIsInstance(out["bbox_regression"], Tensor)
         self.assertIsInstance(out["objectness"], Tensor)
 
-    def test_train_one_epoch(self):
+    def test_training_step(self):
         # Setup the DataModule
-        train_dataset = DummyCOCODetectionDataset(num_samples=128)
+        train_dataset = data_helper.DummyCOCODetectionDataset(num_samples=128)
         data_module = DetectionDataModule(train_dataset, batch_size=16)
         # Load model
         model = yolov5s()
@@ -109,6 +87,43 @@ def test_train_one_epoch(self):
         trainer = pl.Trainer(max_epochs=1)
         trainer.fit(model, data_module)
 
+    def test_vanilla_coco_evaluator(self):
+        # Acquire the images and labels from the coco128 dataset
+        val_dataloader = data_helper.get_dataloader(data_root='data-bin', mode='val')
+        coco = data_helper.get_coco_api_from_dataset(val_dataloader.dataset)
+        coco_evaluator = COCOEvaluator(coco)
+        # Load model
+        model = yolov5s(pretrained=True, score_thresh=0.001)
+        model.eval()
+        for images, targets in val_dataloader:
+            preds = model(images)
+            coco_evaluator.update(preds, targets)
+
+        results = coco_evaluator.compute()
+        self.assertGreater(results['AP'], 41.5)
+        self.assertGreater(results['AP50'], 62.0)
+
+    def tets_test_epoch_end(self):
+        # Acquire the annotation file
+        data_path = Path('data-bin')
+        coco128_dirname = 'coco128'
+        data_helper.prepare_coco128(data_path, dirname=coco128_dirname)
+        annotation_file = data_path / coco128_dirname / 'annotations' / 'instances_train2017.json'
+
+        # Get dataloader to test
+        val_dataloader = data_helper.get_dataloader(data_root=data_path, mode='val')
+
+        # Load model
+        model = yolov5s(pretrained=True, annotation_path=annotation_file)
+        model.eval()
+        # test step
+        trainer = pl.Trainer(max_epochs=1)
+        trainer.test(model, test_dataloaders=val_dataloader)
+        # test epoch end
+        results = model.evaluator.compute()
+        self.assertGreater(results['AP'], 41.5)
+        self.assertGreater(results['AP50'], 62.0)
+
     def test_predict_with_vanilla_model(self):
         # Set image inputs
         img_name = "test/assets/zidane.jpg"

yolort/data/__init__.py

@@ -1,3 +1,4 @@
 # Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
+from .coco_eval import COCOEvaluator
 from .data_pipeline import DataPipeline
-from .data_module import DetectionDataModule, VOCDetectionDataModule, CocoDetectionDataModule
+from .data_module import DetectionDataModule, VOCDetectionDataModule, COCODetectionDataModule