From ae55a54db588cd67c16f8d8eb504fde522a9ef27 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Mon, 28 Jun 2021 03:33:22 -0400 Subject: [PATCH 1/6] Refactor COCODetectionDataModule in data_module.py --- yolort/data/data_module.py | 27 +++++++++++++-------------- yolort/train.py | 5 ++++- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/yolort/data/data_module.py b/yolort/data/data_module.py index dead76a7..8354f557 100644 --- a/yolort/data/data_module.py +++ b/yolort/data/data_module.py @@ -82,8 +82,13 @@ class COCODetectionDataModule(DetectionDataModule): def __init__( self, data_path: str, - annotations_path: Optional[str] = None, - year: str = "2017", + anno_path: Optional[str] = None, + num_classes: int = 80, + data_task: str = "instances", + train_set: str = "train2017", + val_set: str = "val2017", + skip_train_set: bool = False, + skip_val_set: bool = False, train_transform: Optional[Callable] = default_train_transforms, val_transform: Optional[Callable] = default_val_transforms, batch_size: int = 1, @@ -91,23 +96,17 @@ def __init__( *args: Any, **kwargs: Any, ) -> None: - if annotations_path is None: - annotations_path = Path(data_path) / 'annotations' - self.annotations_path = annotations_path + anno_path = Path(anno_path) if anno_path else Path(data_path) / 'annotations' + train_ann_file = anno_path / f"{data_task}_{train_set}.json" + val_ann_file = anno_path / f"{data_task}_{val_set}.json" - train_dataset = self.build_datasets( - data_path, image_set='train', year=year, transforms=train_transform) - val_dataset = self.build_datasets( - data_path, image_set='val', year=year, transforms=val_transform) + train_dataset = None if skip_train_set else COCODetection(data_path, train_ann_file, train_transform()) + val_dataset = None if skip_val_set else COCODetection(data_path, val_ann_file, val_transform()) super().__init__(train_dataset=train_dataset, val_dataset=val_dataset, batch_size=batch_size, num_workers=num_workers, *args, **kwargs) - self.num_classes = 80 - - def build_datasets(self, data_path, image_set, year, transforms): - ann_file = self.annotations_path / f"instances_{image_set}{year}.json" - return COCODetection(data_path, ann_file, transforms()) + self.num_classes = num_classes class VOCDetectionDataModule(DetectionDataModule): diff --git a/yolort/train.py b/yolort/train.py index f49e3d25..05cb083c 100644 --- a/yolort/train.py +++ b/yolort/train.py @@ -41,7 +41,7 @@ def get_args_parser(): def main(args): # Load the data - datamodule = VOCDetectionDataModule.from_argparse_args(args) + datamodule = COCODetectionDataModule.from_argparse_args(args) # Build the model model = models.__dict__[args.arch](num_classes=datamodule.num_classes) @@ -52,6 +52,9 @@ def main(args): # Train the model trainer.fit(model, datamodule=datamodule) + # Save it! + trainer.save_checkpoint("object_detection_model.pt") + if __name__ == "__main__": parser = argparse.ArgumentParser('YOLOv5 training and evaluation script', parents=[get_args_parser()]) From 0d0354e68056e4d5e95991325b7a14cc209e4526 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Mon, 28 Jun 2021 04:41:54 -0400 Subject: [PATCH 2/6] Fixing parameters call in trainer --- yolort/data/data_module.py | 15 +++++++++------ yolort/models/yolo_module.py | 8 ++++---- yolort/train.py | 34 ++++++++++++++++++---------------- 3 files changed, 31 insertions(+), 26 deletions(-) diff --git a/yolort/data/data_module.py b/yolort/data/data_module.py index 8354f557..2e5f60f2 100644 --- a/yolort/data/data_module.py +++ b/yolort/data/data_module.py @@ -22,7 +22,7 @@ def __init__( train_dataset: Optional[Dataset] = None, val_dataset: Optional[Dataset] = None, test_dataset: Optional[Dataset] = None, - batch_size: int = 1, + batch_size: int = 16, num_workers: int = 0, *args: Any, **kwargs: Any, @@ -36,7 +36,7 @@ def __init__( self.batch_size = batch_size self.num_workers = num_workers - def train_dataloader(self, batch_size: int = 16) -> None: + def train_dataloader(self) -> None: """ VOCDetection and COCODetection Args: @@ -44,8 +44,11 @@ def train_dataloader(self, batch_size: int = 16) -> None: transforms: custom transforms """ # Creating data loaders - sampler = torch.utils.data.RandomSampler(self._train_dataset) - batch_sampler = torch.utils.data.BatchSampler(sampler, batch_size, drop_last=True) + batch_sampler = torch.utils.data.BatchSampler( + torch.utils.data.RandomSampler(self._train_dataset), + self.batch_size, + drop_last=True, + ) loader = torch.utils.data.DataLoader( self._train_dataset, @@ -56,7 +59,7 @@ def train_dataloader(self, batch_size: int = 16) -> None: return loader - def val_dataloader(self, batch_size: int = 16) -> None: + def val_dataloader(self) -> None: """ VOCDetection and COCODetection Args: @@ -68,7 +71,7 @@ def val_dataloader(self, batch_size: int = 16) -> None: loader = torch.utils.data.DataLoader( self._val_dataset, - batch_size, + self.batch_size, sampler=sampler, drop_last=False, collate_fn=collate_fn, diff --git a/yolort/models/yolo_module.py b/yolort/models/yolo_module.py index 29f30448..767878f6 100644 --- a/yolort/models/yolo_module.py +++ b/yolort/models/yolo_module.py @@ -25,7 +25,7 @@ class YOLOModule(LightningModule): def __init__( self, lr: float = 0.01, - arch: str = 'yolov5_darknet_pan_s_r31', + arch: str = 'yolov5_darknet_pan_s_r40', pretrained: bool = False, progress: bool = True, size: Tuple[int, int] = (640, 640), @@ -249,14 +249,14 @@ def collate_images(self, samples: Any, image_loader: Callable) -> List[Tensor]: @staticmethod def add_model_specific_args(parent_parser): parser = argparse.ArgumentParser(parents=[parent_parser], add_help=False) - parser.add_argument('--arch', default='yolov5_darknet_pan_s_r31', + parser.add_argument('--arch', default='yolov5_darknet_pan_s_r40', help='model architecture') parser.add_argument('--num_classes', default=80, type=int, help='number classes of datasets') parser.add_argument('--pretrained', action='store_true', help='Use pre-trained models from the modelzoo') - parser.add_argument('--lr', default=0.02, type=float, - help='initial learning rate, 0.02 is the default value for training ' + parser.add_argument('--lr', default=0.01, type=float, + help='initial learning rate, 0.01 is the default value for training ' 'on 8 gpus and 2 images_per_gpu') parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum') diff --git a/yolort/train.py b/yolort/train.py index 05cb083c..713305af 100644 --- a/yolort/train.py +++ b/yolort/train.py @@ -12,28 +12,30 @@ def get_args_parser(): parser.add_argument('--arch', default='yolov5s', help='model structure to train') - parser.add_argument('--data_path', default='./data-bin', - help='dataset') - parser.add_argument('--dataset_type', default='coco', - help='dataset') - parser.add_argument('--dataset_mode', default='instances', - help='dataset mode') - parser.add_argument('--years', default=['2017'], nargs='+', - help='dataset year') - parser.add_argument('--train_set', default='train', - help='set of train') - parser.add_argument('--val_set', default='val', - help='set of val') - parser.add_argument('--batch_size', default=32, type=int, - help='images per gpu, the total batch size is $NGPU x batch_size') parser.add_argument('--max_epochs', default=1, type=int, metavar='N', help='number of total epochs to run') parser.add_argument('--num_gpus', default=1, type=int, metavar='N', help='number of gpu utilizing (default: 1)') + + parser.add_argument('--data_path', default='./data-bin', + help='root path of the dataset') + parser.add_argument('--anno_path', default=None, + help='root path of annotation files') + parser.add_argument('--data_task', default='instances', + help='dataset mode') + parser.add_argument('--train_set', default='train2017', + help='name of train dataset') + parser.add_argument('--val_set', default='val2017', + help='name of val dataset') + parser.add_argument('--skip_train_set', action='store_true', + help='Skip train set') + parser.add_argument('--skip_val_set', action='store_true', + help='Skip val set') + parser.add_argument('--batch_size', default=32, type=int, + help='images per gpu, the total batch size is $NGPU x batch_size') parser.add_argument('--num_workers', default=4, type=int, metavar='N', help='number of data loading workers (default: 4)') - parser.add_argument('--print_freq', default=20, type=int, - help='print frequency') + parser.add_argument('--output_dir', default='.', help='path where to save') return parser From e479ca4bb981d85e964a6a76648bb6c330a1dca2 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Mon, 28 Jun 2021 05:35:46 -0400 Subject: [PATCH 3/6] Fixing arguments --- yolort/models/yolo_module.py | 4 +--- yolort/train.py | 2 ++ 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/yolort/models/yolo_module.py b/yolort/models/yolo_module.py index 767878f6..c09997d3 100644 --- a/yolort/models/yolo_module.py +++ b/yolort/models/yolo_module.py @@ -177,7 +177,7 @@ def configure_optimizers(self): self.model.parameters(), lr=self.lr, momentum=0.9, - weight_decay=0.005, + weight_decay=5e-4, ) @torch.no_grad() @@ -251,8 +251,6 @@ def add_model_specific_args(parent_parser): parser = argparse.ArgumentParser(parents=[parent_parser], add_help=False) parser.add_argument('--arch', default='yolov5_darknet_pan_s_r40', help='model architecture') - parser.add_argument('--num_classes', default=80, type=int, - help='number classes of datasets') parser.add_argument('--pretrained', action='store_true', help='Use pre-trained models from the modelzoo') parser.add_argument('--lr', default=0.01, type=float, diff --git a/yolort/train.py b/yolort/train.py index 713305af..6d78bd15 100644 --- a/yolort/train.py +++ b/yolort/train.py @@ -21,6 +21,8 @@ def get_args_parser(): help='root path of the dataset') parser.add_argument('--anno_path', default=None, help='root path of annotation files') + parser.add_argument('--num_classes', default=80, type=int, + help='number of classes') parser.add_argument('--data_task', default='instances', help='dataset mode') parser.add_argument('--train_set', default='train2017', From 0ad25f674e601e3b9c540e0982a71ffd8a985e34 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Mon, 28 Jun 2021 06:00:30 -0400 Subject: [PATCH 4/6] Port test_data_pipeline.py to pytest --- test/test_data_pipeline.py | 107 +++++++++++++++++++------------------ 1 file changed, 55 insertions(+), 52 deletions(-) diff --git a/test/test_data_pipeline.py b/test/test_data_pipeline.py index cf478fbd..f7864c93 100644 --- a/test/test_data_pipeline.py +++ b/test/test_data_pipeline.py @@ -1,6 +1,6 @@ # Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved. from pathlib import Path -import unittest +import pytest import numpy as np import torch @@ -11,60 +11,63 @@ from typing import Dict -class DataPipelineTester(unittest.TestCase): - def test_contains_any_tensor(self): - dummy_numpy = np.random.randn(3, 6) - self.assertFalse(contains_any_tensor(dummy_numpy)) - dummy_tensor = torch.randn(3, 6) - self.assertTrue(contains_any_tensor(dummy_tensor)) - dummy_tensors = [torch.randn(3, 6), torch.randn(9, 5)] - self.assertTrue(contains_any_tensor(dummy_tensors)) +def test_contains_any_tensor(): + dummy_numpy = np.random.randn(3, 6) + assert not contains_any_tensor(dummy_numpy) + dummy_tensor = torch.randn(3, 6) + assert contains_any_tensor(dummy_tensor) + dummy_tensors = [torch.randn(3, 6), torch.randn(9, 5)] + assert contains_any_tensor(dummy_tensors) - def test_get_dataset(self): - # Acquire the images and labels from the coco128 dataset - train_dataset = data_helper.get_dataset(data_root='data-bin', mode='train') - # Test the datasets - image, target = next(iter(train_dataset)) - self.assertIsInstance(image, Tensor) - self.assertIsInstance(target, Dict) - def test_get_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)) +def test_get_dataset(): + # Acquire the images and labels from the coco128 dataset + train_dataset = data_helper.get_dataset(data_root='data-bin', mode='train') + # Test the datasets + image, target = next(iter(train_dataset)) + assert isinstance(image, Tensor) + assert isinstance(target, Dict) - self.assertEqual(len(images), batch_size) - self.assertIsInstance(images[0], Tensor) - self.assertEqual(len(images[0]), 3) - self.assertEqual(len(targets), batch_size) - self.assertIsInstance(targets[0], Dict) - self.assertIsInstance(targets[0]["image_id"], Tensor) - self.assertIsInstance(targets[0]["boxes"], Tensor) - self.assertIsInstance(targets[0]["labels"], Tensor) - self.assertIsInstance(targets[0]["orig_size"], Tensor) - def test_detection_data_module(self): - # Setup the DataModule - batch_size = 4 - train_dataset = data_helper.get_dataset(data_root='data-bin', mode='train') - data_module = DetectionDataModule(train_dataset, batch_size=batch_size) - self.assertEqual(data_module.batch_size, batch_size) +def test_get_dataloader(): + 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)) - data_loader = data_module.train_dataloader(batch_size=batch_size) - images, targets = next(iter(data_loader)) - self.assertEqual(len(images), batch_size) - self.assertIsInstance(images[0], Tensor) - self.assertEqual(len(images[0]), 3) - self.assertEqual(len(targets), batch_size) - self.assertIsInstance(targets[0], Dict) - self.assertIsInstance(targets[0]["image_id"], Tensor) - self.assertIsInstance(targets[0]["boxes"], Tensor) - self.assertIsInstance(targets[0]["labels"], Tensor) + assert len(images) == batch_size + assert isinstance(images[0], Tensor) + assert len(images[0]) == 3 + assert len(targets) == batch_size + assert isinstance(targets[0], Dict) + assert isinstance(targets[0]["image_id"], Tensor) + assert isinstance(targets[0]["boxes"], Tensor) + assert isinstance(targets[0]["labels"], Tensor) + assert isinstance(targets[0]["orig_size"], Tensor) - def test_prepare_coco128(self): - 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' - self.assertTrue(annotation_file.is_file()) + +def test_detection_data_module(): + # Setup the DataModule + batch_size = 4 + train_dataset = data_helper.get_dataset(data_root='data-bin', mode='train') + data_module = DetectionDataModule(train_dataset, batch_size=batch_size) + assert data_module.batch_size == batch_size + + data_loader = data_module.train_dataloader(batch_size=batch_size) + images, targets = next(iter(data_loader)) + assert len(images) == batch_size + assert isinstance(images[0], Tensor) + assert len(images[0]) == 3 + assert len(targets) == batch_size + assert isinstance(targets[0], Dict) + assert isinstance(targets[0]["image_id"], Tensor) + assert isinstance(targets[0]["boxes"], Tensor) + assert isinstance(targets[0]["labels"], Tensor) + + +def test_prepare_coco128(): + 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' + assert annotation_file.is_file() From b30ce83a767e0bccad285f645eb7913b5e4cd403 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Mon, 28 Jun 2021 06:02:36 -0400 Subject: [PATCH 5/6] Fixing unittest for data_module --- test/test_data_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_data_pipeline.py b/test/test_data_pipeline.py index f7864c93..044c7280 100644 --- a/test/test_data_pipeline.py +++ b/test/test_data_pipeline.py @@ -53,7 +53,7 @@ def test_detection_data_module(): data_module = DetectionDataModule(train_dataset, batch_size=batch_size) assert data_module.batch_size == batch_size - data_loader = data_module.train_dataloader(batch_size=batch_size) + data_loader = data_module.train_dataloader() images, targets = next(iter(data_loader)) assert len(images) == batch_size assert isinstance(images[0], Tensor) From e95e76b8c1efc2e055efb95bb729dfcbf83f2f79 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Mon, 28 Jun 2021 06:16:13 -0400 Subject: [PATCH 6/6] Minor fixes in unittest --- test/common_utils.py | 4 +--- test/test_data_pipeline.py | 1 - test/test_onnx.py | 2 +- 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/test/common_utils.py b/test/common_utils.py index f40aeea0..20a6a7b4 100644 --- a/test/common_utils.py +++ b/test/common_utils.py @@ -17,8 +17,6 @@ import numpy as np from PIL import Image -import logging -logger = logging.getLogger(__name__) def set_rng_seed(seed): torch.manual_seed(seed) @@ -132,7 +130,7 @@ def assertExpected(self, output, subname=None, prec=None, strip_suffix=None): if ACCEPT: filename = {os.path.basename(expected_file)} - logger.info("Accepting updated output for {}:\n\n{}".format(filename, output)) + # logger.info("Accepting updated output for {}:\n\n{}".format(filename, output)) torch.save(output, expected_file) MAX_PICKLE_SIZE = 50 * 1000 # 50 KB binary_size = os.path.getsize(expected_file) diff --git a/test/test_data_pipeline.py b/test/test_data_pipeline.py index 044c7280..ef1ce586 100644 --- a/test/test_data_pipeline.py +++ b/test/test_data_pipeline.py @@ -1,6 +1,5 @@ # Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved. from pathlib import Path -import pytest import numpy as np import torch diff --git a/test/test_onnx.py b/test/test_onnx.py index 20b84c78..cacc3265 100644 --- a/test/test_onnx.py +++ b/test/test_onnx.py @@ -157,7 +157,7 @@ def test_yolotr(self): upstream_version='r4.0', export_friendly=True, pretrained=True, - size=(640, 640), + size=(640, 640), score_thresh=0.45, ) model.eval()