Add Weighted Sampler for highly imbalanced datasets

train.py

@@ -202,6 +202,8 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                                               quad=opt.quad,
                                               prefix=colorstr('train: '),
                                               shuffle=True,
+                                              validation=False,
+                                              weighted_sampler=opt.weighted_sampler,
                                               seed=opt.seed)
     labels = np.concatenate(dataset.labels, 0)
     mlc = int(labels[:, 0].max())  # max label class
@@ -220,7 +222,9 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                                        rank=-1,
                                        workers=workers * 2,
                                        pad=0.5,
-                                       prefix=colorstr('val: '))[0]
+                                       prefix=colorstr('val: '),
+                                       validation=True,
+                                       weighted_sampler=False)[0]
 
         if not resume:
             if not opt.noautoanchor:
@@ -470,6 +474,9 @@ def parse_opt(known=False):
     parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
     parser.add_argument('--seed', type=int, default=0, help='Global training seed')
     parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
+    parser.add_argument('--weighted_sampler',
+                        action='store_true',
+                        help='Use Weighted Sampler (for highly imbalanced data)')
 
     # Logger arguments
     parser.add_argument('--entity', default=None, help='Entity')

utils/dataloaders.py

@@ -26,6 +26,7 @@
 import yaml
 from PIL import ExifTags, Image, ImageOps
 from torch.utils.data import DataLoader, Dataset, dataloader, distributed
+from torch.utils.data.sampler import WeightedRandomSampler
 from tqdm import tqdm
 
 from utils.augmentations import (Albumentations, augment_hsv, classify_albumentations, classify_transforms, copy_paste,
@@ -100,6 +101,47 @@ def seed_worker(worker_id):
     random.seed(worker_seed)
 
 
+def create_weighted_sampler(dataset):
+    labels_per_class = [label[:, 0].tolist() for label in dataset.labels if label.shape[0] > 0]
+    # flatten 2d array into 1d: https://stackoverflow.com/questions/29244286/how-to-flatten-a-2d-list-to-1d-without-using-numpy
+    labels_per_class = [j for sub in labels_per_class for j in sub]
+
+    labels_per_class = np.array(labels_per_class)
+
+    background_count = len([1 for label in dataset.labels if label.shape[0] == 0])
+
+    unique_classes, counts = np.unique(labels_per_class, return_counts=True)
+
+    # = counts / (np.sum(counts) + background_count)
+    # normalized_background = background_count / (np.sum(counts) + background_count)
+
+    weight_cls = 1 / counts
+
+    # create a dictionary for the weight of each class
+    weight_dict = {}
+    for _cls, weight in zip(unique_classes, weight_cls):
+        weight_dict[_cls] = weight
+
+    weight_background = 1 / background_count
+
+    final_weights = []
+    for label in dataset.labels:
+        if label.shape[0] == 0:
+            final_weights.append(weight_background)
+        else:
+            # use weighted sum of labels for weight in case there are multiple labels for the same image
+            label_classes = np.unique(label[:, 0]).tolist()
+            values = []
+            for cls_ in label_classes:
+                values.append(weight_dict[_cls])
+
+            final_weights.append(sum(values) / len(values))
+
+    final_weights = np.array(final_weights)
+    # you can set the num_samples argument to anything. It basically changes your iteration count in every epoch
+    return WeightedRandomSampler(weights=torch.from_numpy(final_weights), num_samples=len(final_weights))
+
+
 def create_dataloader(path,
                       imgsz,
                       batch_size,
@@ -116,6 +158,8 @@ def create_dataloader(path,
                       quad=False,
                       prefix='',
                       shuffle=False,
+                      validation=False,
+                      weighted_sampler=False,
                       seed=0):
     if rect and shuffle:
         LOGGER.warning('WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False')
@@ -138,7 +182,14 @@ def create_dataloader(path,
     batch_size = min(batch_size, len(dataset))
     nd = torch.cuda.device_count()  # number of CUDA devices
     nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers])  # number of workers
+
     sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
+
+    if not validation and weighted_sampler:
+        # weighted sampler should not be called on validation as this will report wrong results
+        assert rank == -1, 'Currently multi-GPU Support is not enabled when using weighted sampler'
+        sampler = create_weighted_sampler(dataset)
+
     loader = DataLoader if image_weights else InfiniteDataLoader  # only DataLoader allows for attribute updates
     generator = torch.Generator()
     generator.manual_seed(6148914691236517205 + seed + RANK)

val.py

@@ -180,7 +180,9 @@ def run(
                                        pad=pad,
                                        rect=rect,
                                        workers=workers,
-                                       prefix=colorstr(f'{task}: '))[0]
+                                       prefix=colorstr(f'{task}: '),
+                                       validation=True,
+                                       weighted_sampler=False)[0]
 
     seen = 0
     confusion_matrix = ConfusionMatrix(nc=nc)