Uneven batch sizes for MultiDatasetDataLoader

[h-0-0]

I've been trying to use the MultiDatasetDataLoader to load in two datasets where I want the batch sizes for each dataset to be different. This is producing the following error:

File "/user/work/jd18380/pyvenv/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py", line 138, in collate raise RuntimeError('each element in list of batch should be of equal size')

RuntimeError: each element in list of batch should be of equal size

This error arises after trying to do the following in the avalanche repository:

File "/user/work/jd18380/pyvenv/lib/python3.10/site-packages/avalanche/benchmarks/utils/data_loader.py", line 191, in iter
yield from self._get_loader()

I managed to fix this problem by commenting out in torch/utils/data/_utils/collate.py:
if not all(len(elem) == elem_size for elem in it): raise RuntimeError('each element in list of batch should be of equal size')
I was wondering if this was an issue that was known? And if it would be worth creating an issue for if not? (as perhaps having to comment out pytorch code isn't the best fix?)

[lrzpellegrini]

Hi @H-aze, can you share a snippet of code we can use to reproduce the issue?

[h-0-0]

Yes of course, here is some code that should generate the error:

from avalanche.benchmarks.utils.data_loader import MultiDatasetDataLoader
from avalanche.core import SupervisedPlugin
from avalanche.training.storage_policy import ExemplarsBuffer
from typing import Any
from avalanche.benchmarks.utils import AvalancheDataset
from torch import arange
from avalanche.benchmarks.classic import SplitCIFAR10
from avalanche.models import SimpleCNN
from torch.optim import Adam
from torch.nn import CrossEntropyLoss
from avalanche.training.templates import SupervisedTemplate
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor

class FixedReplay(SupervisedPlugin):

    def __init__(self, storage_policy, buffer_data, max_size, bs1, bs2):
        """ A simple replay plugin with reservoir sampling. """
        super().__init__()
        self.storage_policy = storage_policy(max_size=max_size)
        self.storage_policy.update_from_dataset(buffer_data)
        self.bs1 = bs1
        self.bs2 = bs2

    def before_training_exp(self, strategy,
                            num_workers: int = 0, shuffle: bool = True,
                            **kwargs):
        """ Here we set the dataloader. We use a dataloader that can load samples from multiple datasets."""
        print("Override the dataloader.")
        strategy.dataloader = MultiDatasetDataLoader(
            datasets = [strategy.adapted_dataset, self.storage_policy.buffer] ,
            batch_sizes = [self.bs1, self.bs2],
            termination_dataset = 0,
            oversample_small_datasets = True,
            distributed_sampling = True,
            shuffle = shuffle,
            num_workers = num_workers,
        )

    def after_training_exp(self, strategy: "BaseStrategy", **kwargs):
        """ We update the buffer after the experience.
            You can use a different callback to update the buffer in a different place
        """
        print("Buffer update.")
        self.storage_policy.update(strategy, **kwargs)

class FixedBuffer(ExemplarsBuffer):
    def __init__(self, max_size: int):
        """
        A buffer for replay that does not modify istelf unless explicitly.
        """
        super().__init__(max_size)

    def update(self, strategy: Any, **kwargs):
        """Update buffer. We wish not to update the buffer automatically."""
        pass

    def update_from_dataset(self, new_data: AvalancheDataset):
        """Replaces buffer with the given dataset.

        :param new_data: the new dataset
        """
        self.buffer = new_data

    def resize(self, strategy: Any, new_size: int):
        """Update the maximum size of the buffer."""
        self.max_size = new_size
        if len(self.buffer) <= self.max_size:
            return
        self.buffer = self.buffer.subset(arange(self.max_size))
        
# GET DATA
scenario = SplitCIFAR10(n_experiences=5, shuffle=False)

# CREATE MODEL
model = SimpleCNN()

# CREATE OPTIMIZER
optimizer = Adam(model.parameters(), lr=0.0001)

buffer_data = CIFAR10(
    "data", 
    download=True, 
    train=True,
    transform=ToTensor()
)

# CREATE THE STRATEGY INSTANCE (NAIVE)
cl_strategy = SupervisedTemplate(
    model, optimizer,
    CrossEntropyLoss(), train_mb_size=128, train_epochs=1, eval_mb_size=128,
    plugins=[FixedReplay(FixedBuffer, buffer_data, max_size=200, bs1=100, bs2=28)]
)

# TRAINING LOOP
print('Starting regular CL experiment...')
train_results = []
test_results = []
for experience in scenario.train_stream:
    print("Start of experience: ", experience.current_experience)
    print("Current Classes: ", experience.classes_in_this_experience)

    # train returns a dictionary which contains all the metric values
    train_results.append(cl_strategy.train(experience))
    print('Training completed')

    print('Computing accuracy on the whole test set')
    # test also returns a dictionary which contains all the metric values
    test_results.append(cl_strategy.eval(scenario.test_stream))
print('Experiment completed')

[lrzpellegrini]

I think I solved it. Change buffer_data to this:

from avalanche.benchmarks.utils.classification_dataset import make_classification_dataset

buffer_data = make_classification_dataset(
    CIFAR10(
        "data", 
        download=True, 
        train=True,
        transform=ToTensor()
    )
)

Also, I noticed that in after_training_exp you are calling self.storage_policy.update(strategy, **kwargs), but the update method is not implemented in FixedBuffer. Maybe you wanted to use update_from_dataset.

Let me know if this solves the issue!

[h-0-0]

Hi @lrzpellegrini, this solves the issue, thanks! I also just found out that adding the following line after creating buffer_data also solves the issue:

buffer_data = [pair + (-1,) for pair in buffer_data]

So I'm guessing its because it needs task labels, which "make_classification_dataset" also does presumably.

[lrzpellegrini]

Hi @lrzpellegrini, this solves the issue, thanks! I also just found out that adding the following line after creating buffer_data also solves the issue:

buffer_data = [pair + (-1,) for pair in buffer_data]

So I'm guessing its because it needs task labels, which "make_classification_dataset" also does presumably.

Yes, that is the reason. I recommend the make_classification_dataset approach as it much faster that concatenating the -1