2603

train.py

@@ -57,7 +57,7 @@ torch_compile = False
 unet_gradient = True
 loss_normalize = False
 fixed_seed = False
-shuffle = False
+shuffle = True
 comet_ml_api_key = "Agctp26mbqnoYrrlvQuKSTk6r" 
 comet_ml_workspace = "recoilme" 
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -247,68 +247,71 @@ if mean is not None and std is not None:
 class DistributedResolutionBatchSampler(Sampler):
     def __init__(self, dataset, batch_size, num_replicas, rank, drop_last=True, shuffle=False):
         self.dataset = dataset
-        self.batch_size = max(1, batch_size // num_replicas)
         self.num_replicas = num_replicas
         self.rank = rank
         self.drop_last = drop_last
         self.shuffle = shuffle
         self.epoch = 0
-        
+
+        # локальный batch для одной GPU
+        self.batch_size = max(1, batch_size // num_replicas)
+        self.global_batch = self.batch_size * self.num_replicas
+
         try:
             widths = np.array(dataset["width"])
             heights = np.array(dataset["height"])
         except KeyError:
             widths = np.zeros(len(dataset))
             heights = np.zeros(len(dataset))
-        
-        # Группируем индексы по размерам
-        self.size_keys = np.unique(np.stack([widths, heights], axis=1), axis=0)
+
+        # группировка индексов по (width, height)
         self.size_groups = {}
-        for w, h in self.size_keys:
-            mask = (widths == w) & (heights == h)
-            self.size_groups[(w, h)] = np.where(mask)[0]
-            
-        self.group_num_batches = {}
-        total_batches = 0
-        for size, indices in self.size_groups.items():
-            num_full_batches = len(indices) // (self.batch_size * self.num_replicas)
-            self.group_num_batches[size] = num_full_batches
-            total_batches += num_full_batches
-            
-        self.num_batches = total_batches
-        
+        for i, (w, h) in enumerate(zip(widths, heights)):
+            self.size_groups.setdefault((w, h), []).append(i)
+
+        # переводим в numpy для скорости
+        for k in self.size_groups:
+            self.size_groups[k] = np.array(self.size_groups[k], dtype=np.int64)
+
+        # считаем общее число батчей
+        self.num_batches = sum(
+            len(indices) // self.global_batch
+            for indices in self.size_groups.values()
+        )
+
     def __iter__(self):
-        all_batches = []
         rng = np.random.RandomState(self.epoch)
-        
-        for size, indices in self.size_groups.items():
-            indices = indices.copy()
+
+        global_batches = []
+
+        for indices in self.size_groups.values():
+            idx = indices.copy()
+
             if self.shuffle:
-                rng.shuffle(indices) # Перемешиваем внутри группы
-            
-            num_full_batches = self.group_num_batches[size]
-            if num_full_batches == 0:
+                rng.shuffle(idx)
+
+            num_batches = len(idx) // self.global_batch
+            if num_batches == 0:
                 continue
-            
-            # Берем только индексы, кратные полному размеру батча (batch_size * num_replicas)
-            valid_indices = indices[:num_full_batches * self.batch_size * self.num_replicas]
-            # Формируем сетку батчей
-            batches = valid_indices.reshape(-1, self.batch_size * self.num_replicas)
-            
-            # Выбираем чанк для конкретного GPU
-            start_idx = self.rank * self.batch_size
-            gpu_batches = batches[:, start_idx : start_idx + self.batch_size]
-            all_batches.extend(gpu_batches.tolist())
 
-        # Перемешиваем батчи разных разрешений между собой
+            idx = idx[:num_batches * self.global_batch]
+            batches = idx.reshape(num_batches, self.global_batch)
+
+            global_batches.extend(batches)
+
         if self.shuffle:
-            rng.shuffle(all_batches)
+            rng.shuffle(global_batches)
 
-        accelerator.wait_for_everyone()
-        return iter(all_batches)
+        # split батча по GPU
+        start = self.rank * self.batch_size
+        end = start + self.batch_size
+
+        result = [batch[start:end] for batch in global_batches]
+
+        return iter(result)
 
     def __len__(self):
-        return self.num_batches 
+        return self.num_batches
 
     def set_epoch(self, epoch):
         self.epoch = epoch