Is there a way to view ARB bucketing results with filenames?

[Aeka0]

Is there any way to see the detailed results after Aspect Ratio Bucketing (ARB)?

The current trainer performs automatic ARB based on the training set resolutions, the target training resolution, and the number of repeats. However, it seems impossible to find out the exact bucketing results (I can only see how many images are in each bucket).

I'm currently optimizing the training process for a small dataset using a high batch size. I've noticed that many buckets end up with very few images. But since I can't determine which specific files are in these sparsely populated buckets, I'm unable to take further action.

Someone told me that these extremely unbalanced buckets can severely impact training performance. I'm hoping someone with similar experience might be able to provide some answers or guidance.

[araleza]

Hi @Aeka0. To answer your question:

Is there any way to see the detailed results after Aspect Ratio Bucketing (ARB)?

...there is now. I've just created a pull request for this feature, as I ran into this issue as well. Having used the feature I've written to fix up my underfilled buckets, my training quality (and resulting image quality) has greatly increased. My change adds debug output like this:

Here's the branch you can grab here if you want the feature now:

#2141

Or, if you don't know how to get pull request branches (or you're training a model that isn't Flux - I've only checked into the SD3 Flux branch, as that's all I've been using), just find this section of code around line 1050 in library/train_util.py:

        # bucket情報を表示、格納する
        if self.enable_bucket:
            self.bucket_info = {"buckets": {}}
            logger.info("number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）")
            for i, (reso, bucket) in enumerate(zip(self.bucket_manager.resos, self.bucket_manager.buckets)):
                count = len(bucket)
                if count > 0:
                    self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)}
                    logger.info(f"bucket {i}: resolution {reso}, count: {len(bucket)}")

            if len(img_ar_errors) == 0:
                mean_img_ar_error = 0  # avoid NaN
            else:
                img_ar_errors = np.array(img_ar_errors)
                mean_img_ar_error = np.mean(np.abs(img_ar_errors))
            self.bucket_info["mean_img_ar_error"] = mean_img_ar_error
            logger.info(f"mean ar error (without repeats): {mean_img_ar_error}")

and add this code snippet in the space in the middle (i.e. before the if len(img_ar_errors) == 0: line):

            # 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}'")

and that should be it working.