Custom Batch Sampler

[josephsuccar]

My workflow involves creating batches that are sampled fairly with respect to some metadata values (x, y, z) such that each batch is guaranteed to sample evenly from x, then evenly from y, then evenly from z. I currently achieve this using the 'batch_sampler' feature of 'torch.utils.data.dataloader' and I'm wondering how I might replicate this using FFCV.

[GuillaumeLeclerc]

Hello,

Does that mean you don't sample each once per epoch right ?

You can look into the TraversalOrder class. That might do what you want.

[josephsuccar]

Thank you for your response.

For metric learning sampling can be done in sets (e.g. pairs, triplets), so the length of epoch might be the permutation of all possible sets, which is absurdly large. For this reason it's typical to use an infinite stream approach while training, and structure code in terms of train steps instead of epochs. Another reason why you might not want to sample one datapoint per epoch is if you have a very uneven class distribution.

However, it seems like the concept of an epoch is tightly coupled with FFCV.Loader. Is it necessary from a design perspective? For example, if the loader needs to be recompiled per epoch, this might be achievable without the loader knowing what an epoch is. Consider an argument replacement recompile (bool) for recompile (int) where the loader can be recompiled after some number of train steps, which could be set equal to an epoch by the user, or not. Also, I'm wondering if TraversalOrder can be instantiated outside of the FFCV.Loader class. Is there anything else that might need to change to have FFCV.Loader reason over iterations instead of epochs?

Thanks again, to you, and the team for open sourcing this great project!

[GuillaumeLeclerc]

Hello, this wouldn't work because when the code is recompiled the iterator unfortunately has to be discarded. Do you have an example of the API you would like and I'll see if I can support that in the upcoming v1.0.

Thank you for the feecback!

[josephsuccar]

I can share an idea of what I'm doing currently.

# toy multi-class example
metadata = pd.DataFrame([{'a': i > 3, 'b': (i+1) % 2, 'c': i % 3, 'id': i} for i in range(13)])

'''
        a  b  c  id
0   False  1  0   0
1   False  0  1   1
2   False  1  2   2
3   False  0  0   3
4    True  1  1   4
5    True  0  2   5
6    True  1  0   6
7    True  0  1   7
8    True  1  2   8
9    True  0  0   9
10   True  1  1  10
11   True  0  2  11
12   True  1  0  12
'''

dataset = NestedPandasDataset(
    metadata=metadata,
    order=['a', 'b'],
    num_samples_per_level=2,
)

batch_sampler=dataset.batch_sampler(batch_size) #   Generator[List[int]], construct batches evenly from dataset classes

dataset = MapDataset(dataset, read_and_augment) #   maybe wrap a func per '__getitem__' of the parent dataset

dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_sampler=batch_sampler,
)

#   TODO test if every example (id) is seen eventually
#   TODO test if mean of samples['a'] is 0.5
#   TODO test if mean of samples['b'] is 0.5

In the above example the loader has no idea how big the dataset is, or what an epoch size should be. That's all handled by the batch sampler. This is good because it accommodates infinite stream style data loading, as well as sampling in sets, where you need some guarantee that a certain number of same or not-same classes appear in each batch.

It would be great for FFCV to have a loader like:

dataloader = FFCV.StreamLoader(
    dataset,
    batch_sampler=batch_sampler,
) # will not Raise StopIteration unless batch_sampler raises StopIteration 

As you mentioned, the most similar thing to a sampler is TraversalOrder. If this could be instantiated outside of the Loader then it would allow for easier customization.

Thank you for your consideration! Hope the above example helps in some way.