fix bug: add sample rate to control different sampling probability for each type of dataset

train_ds.py

@@ -57,6 +57,7 @@ def parse_args(args):
     parser.add_argument(
         "--refer_seg_data", default="refclef||refcoco||refcoco+||refcocog", type=str
     )
+    parser.add_argument("--sample_rates", default="9,3,3,1", type=str)
     parser.add_argument("--vqa_data", default="llava_instruct_150k", type=str)
     parser.add_argument("--reason_seg_data", default="ReasonSeg|train", type=str)
     parser.add_argument("--val_dataset", default="ReasonSeg|val", type=str)
@@ -153,6 +154,7 @@ def main(args):
         num_classes_per_sample=args.num_classes_per_sample,
         exclude_val=args.exclude_val,
         dataset=args.dataset,
+        sample_rate=[float(x) for x in args.sample_rates.split(",")],
         sem_seg_data=args.sem_seg_data,
         refer_seg_data=args.refer_seg_data,
         vqa_data=args.vqa_data,

utils/dataset.py

@@ -152,6 +152,7 @@ class HybridDataset(torch.utils.data.Dataset):
         num_classes_per_sample: int = 3,
         exclude_val=False,
         dataset="sem_seg||refer_seg||vqa||reason_seg",
+        sample_rate=[9, 3, 3, 1],
         sem_seg_data="ade20k||cocostuff||partimagenet||pascal_part||paco_lvis||mapillary",
         refer_seg_data="refclef||refcoco||refcoco+||refcocog",
         vqa_data="llava_instruct_150k",
@@ -163,6 +164,8 @@ class HybridDataset(torch.utils.data.Dataset):
         self.samples_per_epoch = samples_per_epoch
         self.explanatory = explanatory
         self.num_classes_per_sample = num_classes_per_sample
+        sample_rate = np.array(sample_rate)
+        self.sample_rate = sample_rate / sample_rate.sum()
 
         self.base_image_dir = base_image_dir
         self.image_size = image_size
@@ -235,9 +238,7 @@ class HybridDataset(torch.utils.data.Dataset):
         return self.samples_per_epoch
 
     def __getitem__(self, idx):
-        ind = (random.randint(0, 2023) * (idx + 1)) % len(
-            self.datasets
-        )  # random.randint(0, len(self.datasets)-1)
+        ind = np.random.choice(list(range(len(self.datasets))), p=self.sample_rate)
         data = self.all_datasets[ind]
         inference = False
         return *data[0], inference