🔀 [Merge] branch 'DDP_BUGS' of github.com:WongKinYiu/yolov9mit into DDP_BUGS

yolo/tools/data_loader.py

@@ -23,6 +23,20 @@ from yolo.utils.dataset_utils import (
 )
 
 
+def tensorlize(data):
+    # TODO Move Tensorlize to helper
+    img_paths, bboxes = zip(*data)
+    max_box = max(bbox.size(0) for bbox in bboxes)
+    padded_bbox_list = []
+    for bbox in bboxes:
+        padding = torch.full((max_box, 5), -1, dtype=torch.float32)
+        padding[: bbox.size(0)] = bbox
+        padded_bbox_list.append(padding)
+    bboxes = np.stack(padded_bbox_list)
+    img_paths = np.array(img_paths)
+    return img_paths, bboxes
+
+
 class YoloDataset(Dataset):
     def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, phase: str = "train2017"):
         augment_cfg = data_cfg.data_augment
@@ -32,19 +46,8 @@ class YoloDataset(Dataset):
         transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()]
         self.transform = AugmentationComposer(transforms, self.image_size)
         self.transform.get_more_data = self.get_more_data
-        self.img_paths, self.bboxes = self.tensorlize(self.load_data(Path(dataset_cfg.path), phase_name))
-
-    def tensorlize(self, data):
-        img_paths, bboxes = zip(*data)
-        max_box = max(bbox.size(0) for bbox in bboxes)
-        padded_bbox_list = []
-        for bbox in bboxes:
-            padding = torch.full((max_box, 5), -1, dtype=torch.float32)
-            padding[: bbox.size(0)] = bbox
-            padded_bbox_list.append(padding)
-        bboxes = torch.stack(padded_bbox_list)
-        img_paths = np.array(img_paths)
-        return img_paths, bboxes
+        img_paths, bboxes = tensorlize(self.load_data(Path(dataset_cfg.path), phase_name))
+        self.img_paths, self.bboxes = img_paths, bboxes
 
     def load_data(self, dataset_path: Path, phase_name: str):
         """
@@ -147,7 +150,7 @@ class YoloDataset(Dataset):
         img_path, bboxes = self.img_paths[idx], self.bboxes[idx]
         valid_mask = bboxes[:, 0] != -1
         img = Image.open(img_path).convert("RGB")
-        return img, bboxes[valid_mask], img_path
+        return img, torch.from_numpy(bboxes[valid_mask]), img_path
 
     def get_more_data(self, num: int = 1):
         indices = torch.randint(0, len(self), (num,))

yolo/tools/solver.py

@@ -147,7 +147,7 @@ class ModelTrainer:
             self.progress.finish_one_epoch(epoch_loss, epoch_idx=epoch_idx)
 
             mAPs = self.validator.solve(self.validation_dataloader, epoch_idx=epoch_idx)
-            if mAPs is not None and self.good_epoch(mAPs):
+            if self.good_epoch(mAPs):
                 self.save_checkpoint(epoch_idx=epoch_idx)
             # TODO: save model if result are better than before
         self.progress.finish_train()
@@ -256,9 +256,8 @@ class ModelValidator:
 
         with open(self.json_path, "w") as f:
             predict_json = collect_prediction(predict_json, self.progress.local_rank)
-            if self.progress.local_rank != 0:
-                return
-            json.dump(predict_json, f)
+            if self.progress.local_rank == 0:
+                json.dump(predict_json, f)
         if hasattr(self, "coco_gt"):
             self.progress.start_pycocotools()
             result = calculate_ap(self.coco_gt, predict_json)