YOUR REVISION MESSAGE

exps/default/__init__.py

@@ -0,0 +1,3 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.

exps/default/yolov5l.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import yolort.models as models
+
+from yolort.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self):
+        self.model = models.__dict__['yolov5l'](upstream_version="r6.0",)
+        self.model.train()
+        return self.model

exps/default/yolov5m.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import yolort.models as models
+
+from yolort.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self):
+        self.model = models.__dict__['yolov5m'](upstream_version="r6.0",)
+        self.model.train()
+        return self.model

exps/default/yolov5m6.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import yolort.models as models
+
+from yolort.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self):
+        self.model = models.__dict__['yolov5m6'](upstream_version="r6.0",)
+        self.model.train()
+        return self.model

exps/default/yolov5n.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import yolort.models as models
+
+from yolort.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self):
+        self.model = models.__dict__['yolov5n'](upstream_version="r6.0",)
+        self.model.train()
+        return self.model

exps/default/yolov5n6.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import yolort.models as models
+
+from yolort.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self):
+        self.model = models.__dict__['yolov5n6'](upstream_version="r6.0",)
+        self.model.train()
+        return self.model

exps/default/yolov5s.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import yolort.models as models
+
+from yolort.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self):
+        self.model = models.__dict__['yolov5s'](upstream_version="r6.0",)
+        self.model.train()
+        return self.model

exps/default/yolov5s6.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import yolort.models as models
+
+from yolort.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self):
+        self.model = models.__dict__['yolov5s6'](upstream_version="r6.0",)
+        self.model.train()
+        return self.model

exps/default/yolov5ts.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import yolort.models as models
+
+from yolort.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+    def __init__(self):
+        super(Exp, self).__init__()
+        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+    def get_model(self):
+        self.model = models.__dict__['yolov5ts'](upstream_version="r6.0",)
+        self.model.train()
+        return self.model

requirements.txt

@@ -29,4 +29,6 @@ pandas
 # extras --------------------------------------
 # pycocotools on PyPI needs python3.7 as minimal
 # pycocotools>=2.0.2  # corresponds to https://github.com/ppwwyyxx/cocoapi
-thop  # FLOPs computation
+thop    # FLOPs computation
+loguru  # Python logging made (stupidly) simple
+Ninja   # a small build system with a focus on speed

test/test_data_pipeline.py

@@ -3,11 +3,23 @@
 
 import numpy as np
 import pytest
-import torch
+import sys
+sys.path.append("../yolort")
 
+import torch
 from torch import Tensor
-from yolort.data import _helper as data_helper
+from yolort.exp import Exp
+from yolort.data import DataPrefetcher
 from yolort.utils import contains_any_tensor
+from torch import distributed as dist
+
+
+def get_world_size() -> int:
+    if not dist.is_available():
+        return 1
+    if not dist.is_initialized():
+        return 1
+    return dist.get_world_size()
 
 
 def test_contains_any_tensor():
@@ -21,28 +33,32 @@ def test_contains_any_tensor():
 
 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")
+    train_dataset = Exp().get_dataset(data_root="data-bin", mode="train", cache_type=None)
     # Test the datasets
-    image, target = next(iter(train_dataset))
-    assert isinstance(image, Tensor)
-    assert isinstance(target, dict)
+    image, target, _, _ = next(iter(train_dataset))
+    assert image.shape == (3, 640, 640)
+    assert target.shape == (50, 5)
 
 
 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))
+    is_distributed = get_world_size() > 1
+    data_loader = Exp().get_data_loader(
+        batch_size=batch_size,
+        is_distributed=is_distributed,
+        no_aug=False,
+        cache_img=None,
+    )
+    prefetcher = DataPrefetcher(data_loader)
+    images, targets = prefetcher.next()
 
     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)
+    assert isinstance(targets[0], Tensor)
+
+test_get_dataloader()
 
 
 @pytest.mark.skip("Remove Lightning dependency")
@@ -65,11 +81,3 @@ def test_detection_data_module():
     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()