Validation skipped when using custom batch sampler

[6gsn]

🐛 Bug

Hi, I'm trying to use the batch sampler which was implemented as below,

class BalancedBatchSampler(BatchSampler):
    def __init__(
        self, labels, n_classes_in_batch, n_samples_in_batch
    ):
        self.labels = labels
        self.label_set = list(set(self.labels))
        self.label_to_indices = {
            label: list(np.where(np.array(self.labels) == label)[0])
            for label in self.label_set
        }
        for l in self.label_set:
            np.random.shuffle(self.label_to_indices[l])
        self.used_label_indices_count = {label: 0 for label in self.label_set}

        self.n_dataset = len(self.labels)
        self.n_classes_in_batch = n_classes_in_batch
        self.n_samples_in_batch = n_samples_in_batch

        self.batch_size = self.n_samples_in_batch * self.n_classes_in_batch

    def __iter__(self):
        count = 0

        while count + self.batch_size < self.n_dataset:
            sampled_classes = np.random.choice(self.label_set, self.n_classes_in_batch, replace=False)
            indices = []

            for class_ in sampled_classes:
                pad_for_class = self.used_label_indices_count[class_]
                samples_for_class = self.label_to_indices[class_][pad_for_class: pad_for_class + self.n_samples_in_batch]
                indices += samples_for_class
                self.used_label_indices_count[class_] += self.n_samples_in_batch

                if self.used_label_indices_count[class_] + self.n_samples_in_batch > len(self.label_to_indices[class_]):
                    np.random.shuffle(self.label_to_indices[class_])
                    self.used_label_indices_count[class_] = 0

            yield indices
            count += self.batch_size

    def __len__(self):
        # step size in 1 epoch
        return self.n_dataset // self.batch_size

The sampler works well in the training stage of trainer.fit(), however, the validation stage has been skipped.

To Reproduce

I confirmed that the behavior has not occurred when using torch.utils.data.sampler.BatchSampler as below,

# train_dataloader+standard sampler * val_dataloader+standard sampler = validation not skipped
def train_dataloader(self):
    train_sampler = BatchSampler(RandomSampler(self.dataset_train), ...)
    return DataLoader(
        self.dataset_train,
        batch_sampler=train_sampler,
        ...
    )

def val_dataloader(self):
    train_sampler = BatchSampler(RandomSampler(self.dataset_val), ...)
    return DataLoader(
        self.dataset_val,
        batch_sampler=val_sampler,
        ...
    )

# train_dataloader+standard sampler * val_dataloader+custom sampler = validation not skipped
def train_dataloader(self):
    train_sampler = BatchSampler(...)
    return DataLoader(...)

def val_dataloader(self):
    val_sampler = BalancedBatchSampler(...)
    return DataLoader(...)

# train_dataloader+custom sampler * val_dataloader+standard sampler = validation skipped
def train_dataloader(self):
    train_sampler = BalancedBatchSampler(...)
    return DataLoader(...)

def val_dataloader(self):
    val_sampler = BatchSampler(...)
    return DataLoader(...)

# train_dataloader+custom sampler * val_dataloader+custom sampler = validation skipped
def train_dataloader(self):
    train_sampler = BalancedBatchSampler(...)
    return DataLoader(...)

def val_dataloader(self):
    val_sampler = BalancedBatchSampler(...)
    return DataLoader(...)

Environment

• PyTorch Lightning Version (e.g., 1.5.0): 1.6.1
• PyTorch Version (e.g., 1.10): 1.11
• Python version (e.g., 3.9): 3.8.3
• OS (e.g., Linux): Linux
• CUDA/cuDNN version: 10.2
• How you installed PyTorch (conda, pip, source): conda

Additional context

Unfortunately, I could not find the reason why the validation stage has been skipped, and what I'm missing in the implementation.
Is there proper practice to use a custom batch sampler?

cc @justusschock @awaelchli @ninginthecloud @rohitgr7

[otaj]

Hi @6gsn, thank you for the bug report! Can you please provide a full minimal working example that is illustrating the problem?

In the meantime, I think I know what is going on. While predicting, we are wrapping your BatchSampler into our wrapper to keep track of generated batches - however, we expect a custom instance that is inheriting from BatchSampler to take same arguments as stock BatchSampler and generally behave similarly.

While the core functionality of BatchSampler is to generate indices from original dataset that make up one batch (and that is totally what your implementation is doing), in the stock PyTorch, it is assumed that it takes a Sampler which generates the sequence of indices and then just groups those indices according to the batch size (https://pytorch.org/docs/stable/data.html#torch.utils.data.BatchSampler).

And since we are wrapping your BatchSampler we also try to do some more with it and in the process, I believe we break it, because it does not obey this contract of what arguments it takes. So what I expect is actually happening in there is some TypeError somewhere which maybe gets unintentionally silenced and it manifests as your validation loop not happening at all.

So, what can you do? You can try to change your implementation in such a way that it obeys similarly to stock PyTorch, i.e. create a custom Sampler, that is generating a stream of indices and then just wrap it into stock BatchSampler. From the look at your code, it should not be very difficult. However, and that falls on us, please also take a bit of time to create minimal working example with this implementation so we can debug if the error is actually wrongly silenced and correct it.

Oh, and currently, similar restrictions apply to DataLoader (we expect certain arguments, even if they do not make much sense for your custom implementation). Once the PR #12981 is merged, we will not have these restrictions for DataLoader, but we will still have them for BatchSampler. I will take a look if it would be possible to lift these restrictions also from BatchSampler

[6gsn]

@otaj

Thank you for your explanation and advice.
Now I understood clearly how the behavior happened, and what supposed to do.

Hence, I re-implemented the feature using Sampler ( just like your recommendation 👍 ).

class BalancedRandomSamper(RandomSampler):
    def __init__(
        self,
        data_source,
        replacement=False,
        num_samples=None,
        generator=None,
        *args,
        **kwargs,
    ):
        super(BalancedRandomSamper, self).__init__(
            data_source, replacement, num_samples, generator
        )
        ....

    def _balanced_sampling(self):
         ....
    def __iter__(self):
         ....

# in data module
def train_dataloader(self):
    train_sampler = BatchSampler(BalancedRandomSamper(self.dataset_train), ...)

Also, it seems working as I expected!

Thank you so much!

[otaj]

Closing as it appears it is solved

[sseveran]

I am facing the same issue. I am building a custom BatchSampler to be able to bucket examples by sequence length. I was able to make it work by using the approach above. I think however the approach is not optimal since it does not support partial batches with fewer elements than the batch size. I have hacked around it to minimize the impact but it would still be good to support custom batch samplers.

I changed my code have what looked to be the same interface as BatchSampler. I don't quite understand why it does not work. It looks like it should be properly instantiated here.

I did find this that appears to specifically check if the batch_sampler is of type BatchSampler and returns a type error if it is not. From the test cases it looks like maybe it should work if I also have a customer data loader.

[otaj]

Hi @sseveran, thank you for reporting your use case! I will take a look at it on Monday. In the meantime, do you think you can share with me the code for your worked around BatchSampler and also for your custom BatchSampler you tried to use before stumbling on this issue?