diff --git a/test/test_data_pipeline.py b/test/test_data_pipeline.py new file mode 100644 index 00000000..98547c32 --- /dev/null +++ b/test/test_data_pipeline.py @@ -0,0 +1,49 @@ +# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved. +from pathlib import Path +import unittest + +import torch +from torch.utils import data + +from yolort.datasets.coco import CocoDetection +from yolort.datasets.transforms import collate_fn, default_train_transforms +from yolort.utils import prepare_coco128 + +from typing import Dict + + +class DataPipelineTester(unittest.TestCase): + + def test_prepare_coco128(self): + data_path = Path('data-bin') + coco128_dirname = 'coco128' + 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_vanilla_dataloader(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()) + # Test the datasets + image, target = next(iter(dataset)) + self.assertIsInstance(image, torch.Tensor) + self.assertIsInstance(target, Dict) + + batch_size = 4 + sampler = data.RandomSampler(dataset) + batch_sampler = data.BatchSampler(sampler, batch_size, drop_last=True) + loader = data.DataLoader(dataset, batch_sampler=batch_sampler, collate_fn=collate_fn, num_workers=0) + # Test the dataloader + images, targets = next(iter(loader)) + + self.assertEqual(len(images), batch_size) + self.assertEqual(len(targets), batch_size) diff --git a/test/test_engine.py b/test/test_engine.py index ab7e2c82..91773f27 100644 --- a/test/test_engine.py +++ b/test/test_engine.py @@ -1,19 +1,26 @@ +# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved. +from pathlib import Path import unittest import torch +from torch.utils import data +from torchvision.io import read_image + import pytorch_lightning as pl 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.datasets.coco import CocoDetection +from yolort.datasets.transforms import collate_fn, default_train_transforms from yolort.datasets import DetectionDataModule +from yolort.utils import prepare_coco128 + from .dataset_utils import DummyCOCODetectionDataset from typing import Dict -from torchvision.io import read_image - def default_loader(img_name, is_half=False): """ @@ -27,7 +34,7 @@ def default_loader(img_name, is_half=False): class EngineTester(unittest.TestCase): - def test_train(self): + def test_train_with_vanilla_model(self): # Do forward over image img_name = "test/assets/zidane.jpg" img_tensor = default_loader(img_name) @@ -49,6 +56,46 @@ def test_train(self): self.assertIsInstance(out["bbox_regression"], torch.Tensor) self.assertIsInstance(out["objectness"], torch.Tensor) + def test_train_with_vanilla_module(self): + """ + For issue #86: + """ + # 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 = data.RandomSampler(dataset) + batch_sampler = data.BatchSampler(sampler, batch_size, drop_last=True) + loader = data.DataLoader(dataset, batch_sampler=batch_sampler, collate_fn=collate_fn, num_workers=0) + # Sample a pair of images/targets + images, targets = next(iter(loader)) + images = [img.to(device) for img in images] + targets = [{k: v.to(device) for k, v in t.items()} for t in targets] + + # Define the model + model = yolov5s(num_classes=80) + model.train() + + out = model(images, targets) + self.assertIsInstance(out, Dict) + self.assertIsInstance(out["cls_logits"], torch.Tensor) + self.assertIsInstance(out["bbox_regression"], torch.Tensor) + self.assertIsInstance(out["objectness"], torch.Tensor) + + @unittest.skip("Just ignore this.") def test_train_one_step(self): # Load model model = yolov5s() @@ -144,7 +191,3 @@ def test_predict_image_files(self): self.assertIsInstance(predictions[0]["boxes"], torch.Tensor) self.assertIsInstance(predictions[0]["labels"], torch.Tensor) self.assertIsInstance(predictions[0]["scores"], torch.Tensor) - - -if __name__ == '__main__': - unittest.main() diff --git a/test/test_models_utils.py b/test/test_models_utils.py deleted file mode 100644 index e69de29b..00000000 diff --git a/yolort/utils/__init__.py b/yolort/utils/__init__.py index 5007a1d3..330c559a 100644 --- a/yolort/utils/__init__.py +++ b/yolort/utils/__init__.py @@ -1,3 +1,4 @@ from .flash_utils import get_callable_dict from .image_utils import cv2_imshow, get_image_from_url, read_image_to_tensor from .update_module_state import update_module_state_from_ultralytics +from .file_utils import prepare_coco128 diff --git a/yolort/utils/file_utils.py b/yolort/utils/file_utils.py new file mode 100644 index 00000000..b830e2f1 --- /dev/null +++ b/yolort/utils/file_utils.py @@ -0,0 +1,33 @@ +# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved. +from pathlib import PosixPath +from zipfile import ZipFile + +import torch + + +def prepare_coco128( + data_path: PosixPath, + dirname: str = 'coco128', +) -> None: + """ + Prepare coco128 dataset to test. + + Args: + data_path (PosixPath): root path of coco128 dataset. + dirname (str): the directory name of coco128 dataset. Default: 'coco128'. + """ + if not data_path.is_dir(): + print(f'Create a new directory: {data_path}') + data_path.mkdir(parents=True, exist_ok=True) + + zip_path = data_path / 'coco128.zip' + coco128_url = 'https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0/coco128.zip' + if not zip_path.is_file(): + print(f'Downloading coco128 datasets form {coco128_url}') + torch.hub.download_url_to_file(coco128_url, zip_path, hash_prefix='a67d2887') + + coco128_path = data_path / dirname + if not coco128_path.is_dir(): + print(f'Unzipping dataset to {coco128_path}') + with ZipFile(zip_path, 'r') as zip_obj: + zip_obj.extractall(data_path)