From 12da080c654e7ee6ed3f32e1207584378c53e1c7 Mon Sep 17 00:00:00 2001
From: araleza <70412719+araleza@users.noreply.github.com>
Date: Tue, 1 Jul 2025 21:34:46 +0100
Subject: [PATCH] Added warning for buckets without enough images to half fill
 a batch

---
 library/train_util.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/library/train_util.py b/library/train_util.py
index 36d419fd..f5ad204c 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -1060,6 +1060,21 @@ class BaseDataset(torch.utils.data.Dataset):
                     self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}
                     logger.info(f"bucket {i}: resolution {reso}, count: {len(bucket)}")
 
+            # Warn of any particularly small buckets compared to the current batch size
+            already_warned_underfilled_bucket = False
+            for i, bucket in enumerate(self.bucket_manager.buckets):
+                if len(bucket) < (self.batch_size / 2):
+
+                    if not already_warned_underfilled_bucket:
+                        logger.warning(f"These buckets have too few entries to even half fill the batch_size:")
+                        already_warned_underfilled_bucket = True
+
+                    img_filename = bucket[0]
+                    display_filename_len = 55
+                    if len(img_filename) > display_filename_len:
+                        img_filename = f"...{img_filename[-(display_filename_len - 3):]}"
+                    logger.warning(f"  Bucket {self.bucket_manager.resos[i]} e.g. '{img_filename}'")
+
             if len(img_ar_errors) == 0:
                 mean_img_ar_error = 0  # avoid NaN
             else: